Update app.py

app.py

@@ -2,13 +2,15 @@ import os
 import json
 import uuid
 import time
-from datetime import datetime
-from flask import Flask, request, Response, jsonify
+import asyncio
+import socketio
 import socketio
 import requests
 import logging
-from threading import Event
 import re
+from datetime import datetime
+from flask import Flask, request, Response, jsonify
+from threading import Event
 
 app = Flask(__name__)
 
@@ -48,8 +50,7 @@ if proxy_url:
 else:
     transport = None
 
-sio = socketio.Client(http_session=transport, logger=False, engineio_logger=False)
-
+sio = socketio.AsyncClient(http_session=transport, logger=False, engineio_logger=False)
 connect_opts = {'transports': ['websocket', 'polling']}
 
 sio_opts = {
@@ -125,7 +126,7 @@ def root():
     })
 
 @app.route('/ai/v1/messages', methods=['POST'])
-def messages():
+async def messages():
     auth_error = validate_api_key()
     if auth_error:
         return auth_error
@@ -139,20 +140,20 @@ def messages():
         input_tokens = calculate_tokens(previous_messages)
 
         msg_id = str(uuid.uuid4())
-        response_event = Event()
+        response_event = asyncio.Event()
         response_text = []
         total_output_tokens = 0
 
         if not stream:
-            return handle_non_stream(previous_messages, msg_id, model, input_tokens)
+            return await handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
         log_request(request.remote_addr, request.path, 200)
 
-        def generate():
+        async def generate():
             nonlocal total_output_tokens
             start_time = time.time()
             last_activity_time = start_time
-            timeout = max(60, input_tokens / 1000)  # 动态设置超时时间，最少60秒
+            timeout = max(300, input_tokens / 100)  # 动态设置超时时间，最少300秒
 
             def send_event(event_type, data):
                 event = create_event(event_type, data)
@@ -160,9 +161,9 @@ def messages():
                     'event_type': event_type,
                     'data': {'content': event}
                 })
-                yield event
+                return event
 
-            yield from send_event("message_start", {
+            yield send_event("message_start", {
                 "type": "message_start",
                 "message": {
                     "id": msg_id,
@@ -175,10 +176,11 @@ def messages():
                     "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
                 },
             })
-            yield from send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
-            yield from send_event("ping", {"type": "ping"})
+            yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+            yield send_event("ping", {"type": "ping"})
 
-            def on_query_progress(data):
+            @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:
@@ -206,7 +208,8 @@ def messages():
                     })
                     response_event.set()
 
-            def on_connect():
+            @sio.event
+            async def connect():
                 logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
                 emit_data = {
                     "version": "2.9",
@@ -222,7 +225,7 @@ def messages():
                     "prompt_source": "user",
                     "query_source": "home"
                 }
-                sio.emit('perplexity_ask', (previous_messages, emit_data))
+                await sio.emit('perplexity_ask', (previous_messages, emit_data))
                 logger.info("Sent query to Perplexity AI", extra={
                     'event_type': 'query_sent',
                     'data': {
@@ -230,20 +233,23 @@ def messages():
                     }
                 })
 
-            sio.on('connect', on_connect)
-            sio.on('query_progress', on_query_progress)
+            async def heartbeat():
+                while not response_event.is_set():
+                    await sio.emit('ping')
+                    await asyncio.sleep(25)
 
             try:
-                sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+                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 > 10:  # 如果10秒内没有活动，记录警告
-                        logger.warning("No activity for 10 seconds", extra={'event_type': 'inactivity_warning'})
-                    sio.sleep(0.1)
+                    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 from send_event("content_block_delta", {
+                        yield send_event("content_block_delta", {
                             "type": "content_block_delta",
                             "index": 0,
                             "delta": {"type": "text_delta", "text": chunk},
@@ -259,7 +265,7 @@ def messages():
                             'elapsed_time': time.time() - start_time
                         }
                     })
-                    yield from send_event("content_block_delta", {
+                    yield send_event("content_block_delta", {
                         "type": "content_block_delta",
                         "index": 0,
                         "delta": {"type": "text_delta", "text": f"Request timed out after {timeout} seconds"},
@@ -267,22 +273,23 @@ def messages():
             
             except Exception as e:
                 logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
-                yield from send_event("content_block_delta", {
+                yield send_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
                     "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
                 })
             finally:
+                heartbeat_task.cancel()
                 if sio.connected:
-                    sio.disconnect()
+                    await sio.disconnect()
 
-            yield from send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            yield from send_event("message_delta", {
+            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 from send_event("message_stop", {"type": "message_stop"})
+            yield send_event("message_stop", {"type": "message_stop"})
 
         return Response(generate(), content_type='text/event-stream')
 
@@ -406,6 +413,6 @@ 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)
+    app.run(host='0.0.0.0', port=port, debug=True)