Update app.py

app.py

@@ -9,9 +9,7 @@ import requests
 import logging
 import re
 import asyncio
-from concurrent.futures import ThreadPoolExecutor
-
-thread_pool = ThreadPoolExecutor(max_workers=10)
+from functools import partial
 
 app = Flask(__name__)
 
@@ -148,157 +146,159 @@ def messages():
 
         log_request(request.remote_addr, request.path, 200)
 
-        def generate():
-            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},
-                    },
+        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}
                 })
-                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
-                    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',
+                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
+                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
                             }
                         })
-                        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',
+
+                if data.get('final', False):
+                    logger.info("Final response received", extra={
+                        'event_type': 'response_complete',
                         'data': {
-                            'message': previous_messages[:100] + '...' if len(previous_messages) > 100 else previous_messages
+                            '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(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={
-                            'event_type': 'request_timeout',
-                            'data': {
-                                'timeout': timeout,
-                                'input_tokens': input_tokens,
-                                'output_tokens': total_output_tokens,
-                                'elapsed_time': time.time() - start_time
-                            }
-                        })
+            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(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": f"Request timed out after {timeout} seconds"},
+                            "delta": {"type": "text_delta", "text": chunk},
                         })
-                
-                except Exception as e:
-                    logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
+
+                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"Error during socket connection: {str(e)}"},
+                        "delta": {"type": "text_delta", "text": f"Request timed out after {timeout} seconds"},
                     })
-                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},
+            
+            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)}"},
                 })
-                yield send_event("message_stop", {"type": "message_stop"})
+            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 loop.run_until_complete(run_socket_io())
+        def run_async():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                return loop.run_until_complete(run_socket_io())
+            finally:
+                loop.close()
 
-        return Response(generate(), content_type='text/event-stream')
+        return Response(run_async(), content_type='text/event-stream')
 
     except Exception as e:
         logger.error(f"Request error: {str(e)}", exc_info=True)
@@ -306,9 +306,6 @@ def messages():
         return jsonify({"error": str(e)}), 400
 
 def handle_non_stream(previous_messages, msg_id, model, input_tokens):
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-
     async def run_non_stream():
         response_text = []
         total_output_tokens = 0
@@ -405,6 +402,8 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         })
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
     try:
         return loop.run_until_complete(run_non_stream())
     except Exception as e: