Update app.py

app.py

@@ -9,7 +9,9 @@ import requests
 import logging
 import re
 import asyncio
-from threading import Thread
+from concurrent.futures import ThreadPoolExecutor
+
+thread_pool = ThreadPoolExecutor(max_workers=10)
 
 app = Flask(__name__)
 
@@ -140,46 +142,48 @@ def messages():
         input_tokens = calculate_tokens(previous_messages)
 
         msg_id = str(uuid.uuid4())
-        response_event = asyncio.Event()
-        response_text = []
-        total_output_tokens = 0
-
+        
         if not stream:
             return handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
         log_request(request.remote_addr, request.path, 200)
 
         def generate():
-            nonlocal total_output_tokens
-            start_time = time.time()
-            last_activity_time = start_time
-            timeout = max(300, input_tokens / 100)  # 动态设置超时时间，最少300秒
-
-            def send_event(event_type, data):
-                event = create_event(event_type, data)
-                logger.info(f"Sending {event_type} event", extra={
-                    'event_type': event_type,
-                    'data': {'content': event}
-                })
-                return event
-
-            yield send_event("message_start", {
-                "type": "message_start",
-                "message": {
-                    "id": msg_id,
-                    "type": "message",
-                    "role": "assistant",
-                    "model": model,
-                    "content": [],
-                    "stop_reason": None,
-                    "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
-                },
-            })
-            yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
-            yield send_event("ping", {"type": "ping"})
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
 
             async def run_socket_io():
+                response_event = asyncio.Event()
+                response_text = []
+                total_output_tokens = 0
+                start_time = time.time()
+                last_activity_time = start_time
+                timeout = max(300, input_tokens / 100)  # 动态设置超时时间，最少300秒
+
+                def send_event(event_type, data):
+                    event = create_event(event_type, data)
+                    logger.info(f"Sending {event_type} event", extra={
+                        'event_type': event_type,
+                        'data': {'content': event}
+                    })
+                    return event
+
+                yield send_event("message_start", {
+                    "type": "message_start",
+                    "message": {
+                        "id": msg_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "model": model,
+                        "content": [],
+                        "stop_reason": None,
+                        "stop_sequence": None,
+                        "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
+                    },
+                })
+                yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+                yield send_event("ping", {"type": "ping"})
+
                 @sio.event
                 async def query_progress(data):
                     nonlocal total_output_tokens, response_text, last_activity_time
@@ -247,7 +251,14 @@ def messages():
                         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)
+                        await asyncio.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():
                         logger.warning(f"Request timed out after {timeout} seconds", extra={
@@ -259,44 +270,33 @@ def messages():
                                 '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)
-                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},
+                        "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
                     })
+                finally:
+                    heartbeat_task.cancel()
+                    if sio.connected:
+                        await sio.disconnect()
 
-            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"},
+                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"})
 
-            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 loop.run_until_complete(run_socket_io())
 
         return Response(generate(), content_type='text/event-stream')
 
@@ -306,75 +306,71 @@ def messages():
         return jsonify({"error": str(e)}), 400
 
 def handle_non_stream(previous_messages, msg_id, model, input_tokens):
-    try:
-        response_event = asyncio.Event()
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    async def run_non_stream():
         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
-                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
-                            }
-                        })
+        @sio.event
+        async def 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()
-
-            @sio.event
-            async def 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"
+            if data.get('final', False):
+                return True
+
+        @sio.event
+        async def 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"
+            }
+            await sio.emit('perplexity_ask', (previous_messages, emit_data))
+
+        try:
+            await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+            await asyncio.wait_for(sio.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
                 }
-                await sio.emit('perplexity_ask', (previous_messages, emit_data))
-
-            try:
-                await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-                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:
-                if sio.connected:
-                    await sio.disconnect()
-
-        # 在新线程中运行异步Socket.IO
-        thread = Thread(target=lambda: asyncio.run(run_socket_io()))
-        thread.start()
-        thread.join()
+            })
+        finally:
+            if sio.connected:
+                await sio.disconnect()
 
         if not response_text:
             logger.warning(f"No response received (non-stream) after {timeout} seconds", extra={
@@ -409,9 +405,13 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         })
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
+    try:
+        return loop.run_until_complete(run_non_stream())
     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:
+        loop.close()
 
 @app.errorhandler(404)
 def not_found(error):