Update app.py

app.py

@@ -1,16 +1,43 @@
 import os
 import json
 import uuid
-import re
 from datetime import datetime
 from flask import Flask, request, Response, jsonify
 import socketio
 import requests
 import logging
-from threading import Event
+from threading import Event, Timer
+import re
 
 app = Flask(__name__)
 
+# 自定义日志格式化器
+class CustomFormatter(logging.Formatter):
+    def format(self, record):
+        log_data = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "message": self.remove_ansi_escape(record.getMessage()),
+        }
+        if hasattr(record, 'event_type'):
+            log_data['event_type'] = record.event_type
+        if hasattr(record, 'data'):
+            log_data['data'] = record.data
+        return json.dumps(log_data, ensure_ascii=False, indent=2)
+
+    def remove_ansi_escape(self, text):
+        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+        return ansi_escape.sub('', text)
+
+def setup_logging():
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler()
+    handler.setFormatter(CustomFormatter())
+    logger.addHandler(handler)
+
+logger = logging.getLogger(__name__)
+
 # 从环境变量中获取API密钥
 API_KEY = os.environ.get('PPLX_KEY')
 
@@ -46,41 +73,15 @@ sio_opts = {
     }
 }
 
-class CustomFormatter(logging.Formatter):
-    def format(self, record):
-        log_data = {
-            "timestamp": self.formatTime(record, self.datefmt),
-            "level": record.levelname,
-            "message": self.remove_ansi_escape(record.getMessage()),
-        }
-        if hasattr(record, 'event_type'):
-            log_data['event_type'] = record.event_type
-        if hasattr(record, 'data'):
-            log_data['data'] = record.data
-        return json.dumps(log_data, ensure_ascii=False, indent=2)
-
-    def remove_ansi_escape(self, text):
-        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
-        return ansi_escape.sub('', text)
-
-def setup_logging():
-    logger = logging.getLogger()
-    logger.setLevel(logging.INFO)
-    handler = logging.StreamHandler()
-    handler.setFormatter(CustomFormatter())
-    logger.addHandler(handler)
-
-logger = logging.getLogger(__name__)
-
 def log_request(ip, route, status):
     timestamp = datetime.now().isoformat()
-    logger.info(f"Request received", extra={
-        'event_type': 'request',
+    logger.info(f"Request logged", extra={
+        'event_type': 'request_log',
         'data': {
-            'timestamp': timestamp,
             'ip': ip,
             'route': route,
-            'status': status
+            'status': status,
+            'timestamp': timestamp
         }
     })
 
@@ -151,6 +152,7 @@ def messages():
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
+        timeout_event = Event()
         response_text = []
         total_output_tokens = 0
 
@@ -241,10 +243,18 @@ def messages():
             sio.on('connect', on_connect)
             sio.on('query_progress', on_query_progress)
 
+            def timeout_handler():
+                logger.warning("Request timed out", extra={'event_type': 'request_timeout'})
+                timeout_event.set()
+                response_event.set()
+
+            timer = Timer(30, timeout_handler)  # 30 seconds timeout
+            timer.start()
+
             try:
                 sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
                 
-                while not response_event.is_set():
+                while not response_event.is_set() and not timeout_event.is_set():
                     sio.sleep(0.1)
                     while response_text:
                         chunk = response_text.pop(0)
@@ -254,6 +264,13 @@ def messages():
                             "delta": {"type": "text_delta", "text": chunk},
                         })
                 
+                if timeout_event.is_set():
+                    yield from send_event("content_block_delta", {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": "Request timed out"},
+                    })
+                
             except Exception as e:
                 logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
                 yield from send_event("content_block_delta", {
@@ -262,6 +279,7 @@ def messages():
                     "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
                 })
             finally:
+                timer.cancel()
                 if sio.connected:
                     sio.disconnect()
 
@@ -288,7 +306,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         total_output_tokens = 0
 
         def on_query_progress(data):
-            nonlocal response_text, total_output_tokens
+            nonlocal total_output_tokens, response_text
             if 'text' in data:
                 text = json.loads(data['text'])
                 chunk = text['chunks'][-1] if text['chunks'] else None
@@ -323,8 +341,13 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
 
         sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
         
+        # Wait for response with timeout
         response_event.wait(timeout=30)
 
+        if not response_text:
+            logger.warning("No response received (non-stream)", extra={'event_type': 'no_response_non_stream'})
+            return jsonify({"error": "No response received"}), 504
+
         full_response = {
             "content": [{"text": ''.join(response_text), "type": "text"}],
             "id": msg_id,
@@ -345,10 +368,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
     except Exception as e:
-        logger.error(f"Error during socket connection: {str(e)}", extra={
-            'event_type': 'connection_error_non_stream',
-            'data': {'error': str(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:
@@ -361,17 +381,14 @@ def not_found(error):
 
 @app.errorhandler(500)
 def server_error(error):
-    logger.error(f"Server error: {str(error)}", extra={
-        'event_type': 'server_error',
-        'data': {'error': str(error)}
-    })
+    logger.error(f"Server error: {str(error)}", exc_info=True)
     log_request(request.remote_addr, request.path, 500)
     return "Something broke!", 500
 
 if __name__ == '__main__':
     setup_logging()
     port = int(os.environ.get('PORT', 8081))
-    logger.info("Perplexity proxy starting", extra={
+    logger.info(f"Perplexity proxy starting", extra={
         'event_type': 'server_start',
         'data': {'port': port}
     })