Update app.py

app.py

@@ -11,6 +11,7 @@ import re
 
 app = Flask(__name__)
 
+# 自定义日志格式化器
 class CustomFormatter(logging.Formatter):
     def format(self, record):
         log_data = {
@@ -37,9 +38,13 @@ def setup_logging():
 
 logger = logging.getLogger(__name__)
 
+# 从环境变量中获取API密钥
 API_KEY = os.environ.get('PPLX_KEY')
+
+# 代理设置
 proxy_url = os.environ.get('PROXY_URL')
 
+# 设置代理
 if proxy_url:
     proxies = {
         'http': proxy_url,
@@ -52,10 +57,12 @@ else:
 
 sio = socketio.Client(http_session=transport, logger=False, engineio_logger=False)
 
+# 连接选项
 connect_opts = {
     'transports': ['websocket', 'polling'],
 }
 
+# 其他选项
 sio_opts = {
     'extraHeaders': {
         'Cookie': os.environ.get('PPLX_COOKIE'),
@@ -107,11 +114,6 @@ def create_event(event, data):
         data = json.dumps(data, ensure_ascii=False)
     return f"event: {event}\ndata: {data}\n\n"
 
-def calculate_timeout(input_tokens):
-    base_timeout = 30
-    additional_time = min(max(input_tokens // 100 * 5, 0), 270)
-    return base_timeout + additional_time
-
 @app.route('/')
 def root():
     log_request(request.remote_addr, request.path, 200)
@@ -127,8 +129,8 @@ def root():
                 },
                 "body": {
                     "messages": "Array of message objects",
-                    "stream": "Boolean (optional, defaults to false)",
-                    "model": "Model to be used (optional, defaults to claude-3-5-sonnet-20240620)"
+                    "stream": "Boolean (true for streaming response)",
+                    "model": "Model to be used (optional, defaults to claude-3-opus-20240229)"
                 }
             }
         }
@@ -147,7 +149,6 @@ def messages():
 
         previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
         input_tokens = calculate_tokens(previous_messages)
-        timeout = calculate_timeout(input_tokens)
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
@@ -156,7 +157,7 @@ def messages():
         total_output_tokens = 0
 
         if not stream:
-            return handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout)
+            return handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
         log_request(request.remote_addr, request.path, 200)
 
@@ -171,6 +172,7 @@ def messages():
                 })
                 yield event
 
+            # Send initial events
             yield from send_event("message_start", {
                 "type": "message_start",
                 "message": {
@@ -242,17 +244,11 @@ def messages():
             sio.on('query_progress', on_query_progress)
 
             def timeout_handler():
-                logger.warning("Request timed out", extra={
-                    'event_type': 'request_timeout',
-                    'data': {
-                        'input_tokens': input_tokens,
-                        'timeout': timeout
-                    }
-                })
+                logger.warning("Request timed out", extra={'event_type': 'request_timeout'})
                 timeout_event.set()
                 response_event.set()
 
-            timer = Timer(timeout, timeout_handler)
+            timer = Timer(60, timeout_handler)  # 60 seconds timeout
             timer.start()
 
             try:
@@ -287,6 +283,7 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
+            # Send final events
             yield from send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
             yield from send_event("message_delta", {
                 "type": "message_delta",
@@ -302,7 +299,7 @@ def messages():
         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, timeout):
+def handle_non_stream(previous_messages, msg_id, model, input_tokens):
     try:
         response_event = Event()
         response_text = []
@@ -344,16 +341,11 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout):
 
         sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
         
-        response_event.wait(timeout=timeout)
+        # 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',
-                'data': {
-                    'input_tokens': input_tokens,
-                    'timeout': timeout
-                }
-            })
+            logger.warning("No response received (non-stream)", extra={'event_type': 'no_response_non_stream'})
             return jsonify({"error": "No response received"}), 504
 
         full_response = {