Update app.py

app.py

@@ -2,15 +2,14 @@ import os
 import json
 import uuid
 import time
-import asyncio
-import socketio
+from datetime import datetime
+from flask import Flask, request, Response, jsonify
 import socketio
 import requests
 import logging
 import re
-from datetime import datetime
-from flask import Flask, request, Response, jsonify
-from threading import Event
+import asyncio
+from threading import Thread
 
 app = Flask(__name__)
 
@@ -51,6 +50,7 @@ else:
     transport = None
 
 sio = socketio.AsyncClient(http_session=transport, logger=False, engineio_logger=False)
+
 connect_opts = {'transports': ['websocket', 'polling']}
 
 sio_opts = {
@@ -126,7 +126,7 @@ def root():
     })
 
 @app.route('/ai/v1/messages', methods=['POST'])
-async def messages():
+def messages():
     auth_error = validate_api_key()
     if auth_error:
         return auth_error
@@ -145,11 +145,11 @@ async def messages():
         total_output_tokens = 0
 
         if not stream:
-            return await handle_non_stream(previous_messages, msg_id, model, input_tokens)
+            return handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
         log_request(request.remote_addr, request.path, 200)
 
-        async def generate():
+        def generate():
             nonlocal total_output_tokens
             start_time = time.time()
             last_activity_time = start_time
@@ -179,10 +179,144 @@ async def messages():
             yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
             yield send_event("ping", {"type": "ping"})
 
+            async def run_socket_io():
+                @sio.event
+                async def query_progress(data):
+                    nonlocal total_output_tokens, response_text, last_activity_time
+                    last_activity_time = time.time()
+                    if 'text' in data:
+                        text = json.loads(data['text'])
+                        chunk = text['chunks'][-1] if text['chunks'] else None
+                        if chunk:
+                            response_text.append(chunk)
+                            chunk_tokens = calculate_tokens(chunk)
+                            total_output_tokens += chunk_tokens
+                            logger.info("Received chunk", extra={
+                                'event_type': 'chunk_received',
+                                'data': {
+                                    'chunk': chunk,
+                                    'tokens': chunk_tokens,
+                                    'total_tokens': total_output_tokens
+                                }
+                            })
+
+                    if data.get('final', False):
+                        logger.info("Final response received", extra={
+                            'event_type': 'response_complete',
+                            'data': {
+                                'total_tokens': total_output_tokens
+                            }
+                        })
+                        response_event.set()
+
+                @sio.event
+                async def connect():
+                    logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
+                    emit_data = {
+                        "version": "2.9",
+                        "source": "default",
+                        "attachments": [],
+                        "language": "en-GB",
+                        "timezone": "Europe/London",
+                        "mode": "concise",
+                        "is_related_query": False,
+                        "is_default_related_query": False,
+                        "visitor_id": str(uuid.uuid4()),
+                        "frontend_context_uuid": str(uuid.uuid4()),
+                        "prompt_source": "user",
+                        "query_source": "home"
+                    }
+                    await sio.emit('perplexity_ask', (previous_messages, emit_data))
+                    logger.info("Sent query to Perplexity AI", extra={
+                        'event_type': 'query_sent',
+                        'data': {
+                            'message': previous_messages[:100] + '...' if len(previous_messages) > 100 else previous_messages
+                        }
+                    })
+
+                async def heartbeat():
+                    while not response_event.is_set():
+                        await sio.emit('ping')
+                        await asyncio.sleep(25)
+
+                try:
+                    await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+                    heartbeat_task = asyncio.create_task(heartbeat())
+                    
+                    while not response_event.is_set() and (time.time() - start_time) < timeout:
+                        current_time = time.time()
+                        if current_time - last_activity_time > 60:  # 如果60秒内没有活动，记录警告
+                            logger.warning("No activity for 60 seconds", extra={'event_type': 'inactivity_warning'})
+                        await asyncio.sleep(1)
+
+                    if not response_event.is_set():
+                        logger.warning(f"Request timed out after {timeout} seconds", extra={
+                            'event_type': 'request_timeout',
+                            'data': {
+                                'timeout': timeout,
+                                'input_tokens': input_tokens,
+                                'output_tokens': total_output_tokens,
+                                'elapsed_time': time.time() - start_time
+                            }
+                        })
+                
+                except Exception as e:
+                    logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
+                finally:
+                    heartbeat_task.cancel()
+                    if sio.connected:
+                        await sio.disconnect()
+
+            # 在新线程中运行异步Socket.IO
+            thread = Thread(target=lambda: asyncio.run(run_socket_io()))
+            thread.start()
+
+            while not response_event.is_set() and (time.time() - start_time) < timeout:
+                time.sleep(0.1)
+                while response_text:
+                    chunk = response_text.pop(0)
+                    yield send_event("content_block_delta", {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": chunk},
+                    })
+
+            if not response_event.is_set():
+                yield send_event("content_block_delta", {
+                    "type": "content_block_delta",
+                    "index": 0,
+                    "delta": {"type": "text_delta", "text": f"Request timed out after {timeout} seconds"},
+                })
+
+            thread.join()  # 等待Socket.IO线程完成
+
+            yield send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+            yield send_event("message_delta", {
+                "type": "message_delta",
+                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                "usage": {"output_tokens": total_output_tokens},
+            })
+            yield send_event("message_stop", {"type": "message_stop"})
+
+        return Response(generate(), content_type='text/event-stream')
+
+    except Exception as e:
+        logger.error(f"Request error: {str(e)}", exc_info=True)
+        log_request(request.remote_addr, request.path, 400)
+        return jsonify({"error": str(e)}), 400
+
+def handle_non_stream(previous_messages, msg_id, model, input_tokens):
+    try:
+        response_event = asyncio.Event()
+        response_text = []
+        total_output_tokens = 0
+        start_time = time.time()
+        timeout = max(300, input_tokens / 100)  # 动态设置超时时间，最少300秒
+
+        async def run_socket_io():
             @sio.event
             async def query_progress(data):
-                nonlocal total_output_tokens, response_text, last_activity_time
-                last_activity_time = time.time()
+                nonlocal total_output_tokens, response_text
                 if 'text' in data:
                     text = json.loads(data['text'])
                     chunk = text['chunks'][-1] if text['chunks'] else None
@@ -190,8 +324,8 @@ async def messages():
                         response_text.append(chunk)
                         chunk_tokens = calculate_tokens(chunk)
                         total_output_tokens += chunk_tokens
-                        logger.info("Received chunk", extra={
-                            'event_type': 'chunk_received',
+                        logger.info("Received chunk (non-stream)", extra={
+                            'event_type': 'chunk_received_non_stream',
                             'data': {
                                 'chunk': chunk,
                                 'tokens': chunk_tokens,
@@ -200,17 +334,11 @@ async def messages():
                         })
 
                 if data.get('final', False):
-                    logger.info("Final response received", extra={
-                        'event_type': 'response_complete',
-                        'data': {
-                            'total_tokens': total_output_tokens
-                        }
-                    })
                     response_event.set()
 
             @sio.event
             async def connect():
-                logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
+                logger.info("Connected to Perplexity AI (non-stream)", extra={'event_type': 'connection_established_non_stream'})
                 emit_data = {
                     "version": "2.9",
                     "source": "default",
@@ -226,132 +354,27 @@ async def messages():
                     "query_source": "home"
                 }
                 await sio.emit('perplexity_ask', (previous_messages, emit_data))
-                logger.info("Sent query to Perplexity AI", extra={
-                    'event_type': 'query_sent',
-                    'data': {
-                        'message': previous_messages[:100] + '...' if len(previous_messages) > 100 else previous_messages
-                    }
-                })
-
-            async def heartbeat():
-                while not response_event.is_set():
-                    await sio.emit('ping')
-                    await asyncio.sleep(25)
 
             try:
                 await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-                heartbeat_task = asyncio.create_task(heartbeat())
-                
-                while not response_event.is_set() and (time.time() - start_time) < timeout:
-                    current_time = time.time()
-                    if current_time - last_activity_time > 60:  # 如果60秒内没有活动，记录警告
-                        logger.warning("No activity for 60 seconds", extra={'event_type': 'inactivity_warning'})
-                    await asyncio.sleep(1)
-                    while response_text:
-                        chunk = response_text.pop(0)
-                        yield send_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
-                
-                if not response_event.is_set():
-                    logger.warning(f"Request timed out after {timeout} seconds", extra={
-                        'event_type': 'request_timeout',
-                        'data': {
-                            'timeout': timeout,
-                            'input_tokens': input_tokens,
-                            'output_tokens': total_output_tokens,
-                            'elapsed_time': time.time() - start_time
-                        }
-                    })
-                    yield send_event("content_block_delta", {
-                        "type": "content_block_delta",
-                        "index": 0,
-                        "delta": {"type": "text_delta", "text": f"Request timed out after {timeout} seconds"},
-                    })
-            
-            except Exception as e:
-                logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
-                yield send_event("content_block_delta", {
-                    "type": "content_block_delta",
-                    "index": 0,
-                    "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
+                await asyncio.wait_for(response_event.wait(), timeout=timeout)
+            except asyncio.TimeoutError:
+                logger.warning(f"Request timed out after {timeout} seconds (non-stream)", extra={
+                    'event_type': 'request_timeout_non_stream',
+                    'data': {
+                        'timeout': timeout,
+                        'input_tokens': input_tokens,
+                        'elapsed_time': time.time() - start_time
+                    }
                 })
             finally:
-                heartbeat_task.cancel()
                 if sio.connected:
                     await sio.disconnect()
 
-            yield send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            yield send_event("message_delta", {
-                "type": "message_delta",
-                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                "usage": {"output_tokens": total_output_tokens},
-            })
-            yield send_event("message_stop", {"type": "message_stop"})
-
-        return Response(generate(), content_type='text/event-stream')
-
-    except Exception as e:
-        logger.error(f"Request error: {str(e)}", exc_info=True)
-        log_request(request.remote_addr, request.path, 400)
-        return jsonify({"error": str(e)}), 400
-
-def handle_non_stream(previous_messages, msg_id, model, input_tokens):
-    try:
-        response_event = Event()
-        response_text = []
-        total_output_tokens = 0
-        start_time = time.time()
-        timeout = max(60, input_tokens / 1000)  # 动态设置超时时间，最少60秒
-
-        def on_query_progress(data):
-            nonlocal total_output_tokens, response_text
-            if 'text' in data:
-                text = json.loads(data['text'])
-                chunk = text['chunks'][-1] if text['chunks'] else None
-                if chunk:
-                    response_text.append(chunk)
-                    chunk_tokens = calculate_tokens(chunk)
-                    total_output_tokens += chunk_tokens
-                    logger.info("Received chunk (non-stream)", extra={
-                        'event_type': 'chunk_received_non_stream',
-                        'data': {
-                            'chunk': chunk,
-                            'tokens': chunk_tokens,
-                            'total_tokens': total_output_tokens
-                        }
-                    })
-
-            if data.get('final', False):
-                response_event.set()
-
-        def on_connect():
-            logger.info("Connected to Perplexity AI (non-stream)", extra={'event_type': 'connection_established_non_stream'})
-            emit_data = {
-                "version": "2.9",
-                "source": "default",
-                "attachments": [],
-                "language": "en-GB",
-                "timezone": "Europe/London",
-                "mode": "concise",
-                "is_related_query": False,
-                "is_default_related_query": False,
-                "visitor_id": str(uuid.uuid4()),
-                "frontend_context_uuid": str(uuid.uuid4()),
-                "prompt_source": "user",
-                "query_source": "home"
-            }
-            sio.emit('perplexity_ask', (previous_messages, emit_data))
-
-        sio.on('connect', on_connect)
-        sio.on('query_progress', on_query_progress)
-
-        sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-        
-        # Wait for response with timeout
-        response_event.wait(timeout=timeout)
+        # 在新线程中运行异步Socket.IO
+        thread = Thread(target=lambda: asyncio.run(run_socket_io()))
+        thread.start()
+        thread.join()
 
         if not response_text:
             logger.warning(f"No response received (non-stream) after {timeout} seconds", extra={
@@ -389,9 +412,6 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
     except Exception as e:
         logger.error(f"Error during non-stream socket connection: {str(e)}", exc_info=True)
         return jsonify({"error": str(e)}), 500
-    finally:
-        if sio.connected:
-            sio.disconnect()
 
 @app.errorhandler(404)
 def not_found(error):
@@ -413,6 +433,4 @@ if __name__ == '__main__':
     })
     if not API_KEY:
         logger.warning("PPLX_KEY environment variable is not set", extra={'event_type': 'config_warning'})
-    app.run(host='0.0.0.0', port=port, debug=True)
-
-
+    app.run(host='0.0.0.0', port=port)