Update app.py

app.py

@@ -1,17 +1,17 @@
 import os
 import json
 import uuid
+import time
 from datetime import datetime
 from flask import Flask, request, Response, jsonify
 import socketio
 import requests
 import logging
-from threading import Event, Timer
+from threading import Event
 import re
 
 app = Flask(__name__)
 
-# 自定义日志格式化器
 class CustomFormatter(logging.Formatter):
     def format(self, record):
         log_data = {
@@ -38,18 +38,11 @@ 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,
-        'https': proxy_url
-    }
+    proxies = {'http': proxy_url, 'https': proxy_url}
     transport = requests.Session()
     transport.proxies.update(proxies)
 else:
@@ -57,12 +50,8 @@ else:
 
 sio = socketio.Client(http_session=transport, logger=False, engineio_logger=False)
 
-# 连接选项
-connect_opts = {
-    'transports': ['websocket', 'polling'],
-}
+connect_opts = {'transports': ['websocket', 'polling']}
 
-# 其他选项
 sio_opts = {
     'extraHeaders': {
         'Cookie': os.environ.get('PPLX_COOKIE'),
@@ -106,8 +95,7 @@ def calculate_tokens(text):
     if re.search(r'[^\x00-\x7F]', text):
         return len(text)
     else:
-        tokens = text.split()
-        return len(tokens)
+        return len(text.split())
 
 def create_event(event, data):
     if isinstance(data, dict):
@@ -129,8 +117,8 @@ def root():
                 },
                 "body": {
                     "messages": "Array of message objects",
-                    "stream": "Boolean (true for streaming response)",
-                    "model": "Model to be used (optional, defaults to claude-3-opus-20240229)"
+                    "stream": "Boolean (optional, defaults to false)",
+                    "model": "Model to be used (optional, defaults to claude-3-5-sonnet-20240620)"
                 }
             }
         }
@@ -152,7 +140,6 @@ def messages():
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
-        timeout_event = Event()
         response_text = []
         total_output_tokens = 0
 
@@ -163,6 +150,9 @@ def messages():
 
         def generate():
             nonlocal total_output_tokens
+            start_time = time.time()
+            last_activity_time = start_time
+            timeout = max(60, input_tokens / 1000)  # 动态设置超时时间，最少60秒
 
             def send_event(event_type, data):
                 event = create_event(event_type, data)
@@ -172,7 +162,6 @@ def messages():
                 })
                 yield event
 
-            # Send initial events
             yield from send_event("message_start", {
                 "type": "message_start",
                 "message": {
@@ -190,7 +179,8 @@ def messages():
             yield from send_event("ping", {"type": "ping"})
 
             def on_query_progress(data):
-                nonlocal total_output_tokens, response_text
+                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
@@ -243,18 +233,13 @@ 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() and not timeout_event.is_set():
+                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)
                     while response_text:
                         chunk = response_text.pop(0)
@@ -264,13 +249,22 @@ def messages():
                             "delta": {"type": "text_delta", "text": chunk},
                         })
                 
-                if timeout_event.is_set():
+                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 from send_event("content_block_delta", {
                         "type": "content_block_delta",
                         "index": 0,
-                        "delta": {"type": "text_delta", "text": "Request timed out"},
+                        "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)
                 yield from send_event("content_block_delta", {
@@ -279,11 +273,9 @@ def messages():
                     "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
                 })
             finally:
-                timer.cancel()
                 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",
@@ -304,6 +296,8 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         response_event = Event()
         response_text = []
         total_output_tokens = 0
+        start_time = time.time()
+        timeout = max(60, input_tokens / 1000)  # 动态设置超时时间，最少60秒
 
         def on_query_progress(data):
             nonlocal total_output_tokens, response_text
@@ -314,6 +308,14 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
                     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()
@@ -342,11 +344,18 @@ 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)
+        response_event.wait(timeout=timeout)
 
         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
+            logger.warning(f"No response received (non-stream) after {timeout} seconds", extra={
+                'event_type': 'no_response_non_stream',
+                'data': {
+                    'timeout': timeout,
+                    'input_tokens': input_tokens,
+                    'elapsed_time': time.time() - start_time
+                }
+            })
+            return jsonify({"error": f"No response received after {timeout} seconds"}), 504
 
         full_response = {
             "content": [{"text": ''.join(response_text), "type": "text"}],
@@ -363,11 +372,14 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
         }
         logger.info("Sending non-stream response", extra={
             'event_type': 'non_stream_response',
-            'data': {'content': full_response}
+            'data': {
+                'content': full_response,
+                'elapsed_time': time.time() - start_time
+            }
         })
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
-    except Exception as e:
+        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:
@@ -394,4 +406,5 @@ 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)
+