Update app.py

app.py

@@ -8,11 +8,9 @@ import requests
 import logging
 from threading import Event, Timer
 import re
-import asyncio
 
 app = Flask(__name__)
 
-# 自定义日志格式化器
 class CustomFormatter(logging.Formatter):
     def format(self, record):
         log_data = {
@@ -39,13 +37,9 @@ 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,
@@ -58,12 +52,10 @@ 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'),
@@ -115,6 +107,11 @@ 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)
@@ -130,18 +127,13 @@ 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)"
                 }
             }
         }
     })
 
-async def process_large_request(previous_messages, model, input_tokens):
-    # 这个函数用于异步处理大型请求
-    # 实际实现时，你可能需要将这个过程放到后台任务队列中
-    pass
-
 @app.route('/ai/v1/messages', methods=['POST'])
 def messages():
     auth_error = validate_api_key()
@@ -150,14 +142,12 @@ def messages():
 
     try:
         json_body = request.json
-        model = json_body.get('model', 'claude-3-opus-20240229')
-        stream = json_body.get('stream', True)
+        model = json_body.get('model', 'claude-3-5-sonnet-20240620')
+        stream = json_body.get('stream', False)
 
         previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
         input_tokens = calculate_tokens(previous_messages)
-
-        # 根据 input_tokens 动态调整超时时间
-        timeout = max(30, min(300, input_tokens // 1000 * 30))  # 最少30秒，最多300秒
+        timeout = calculate_timeout(input_tokens)
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
@@ -165,14 +155,6 @@ def messages():
         response_text = []
         total_output_tokens = 0
 
-        if input_tokens > 100000:  # 如果 tokens 数量特别大，使用异步处理
-            task_id = str(uuid.uuid4())
-            asyncio.create_task(process_large_request(previous_messages, model, input_tokens))
-            return jsonify({
-                "message": "Request is being processed asynchronously",
-                "task_id": task_id
-            }), 202
-
         if not stream:
             return handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout)
 
@@ -189,7 +171,6 @@ def messages():
                 })
                 yield event
 
-            # Send initial events
             yield from send_event("message_start", {
                 "type": "message_start",
                 "message": {
@@ -223,12 +204,6 @@ def messages():
                                 'total_tokens': total_output_tokens
                             }
                         })
-                        # 发送进度更新
-                        yield from send_event("progress", {
-                            "type": "progress",
-                            "processed_tokens": total_output_tokens,
-                            "total_tokens": input_tokens
-                        })
 
                 if data.get('final', False):
                     logger.info("Final response received", extra={
@@ -267,11 +242,17 @@ def messages():
             sio.on('query_progress', on_query_progress)
 
             def timeout_handler():
-                logger.warning("Request timed out", extra={'event_type': 'request_timeout'})
+                logger.warning("Request timed out", extra={
+                    'event_type': 'request_timeout',
+                    'data': {
+                        'input_tokens': input_tokens,
+                        'timeout': timeout
+                    }
+                })
                 timeout_event.set()
                 response_event.set()
 
-            timer = Timer(timeout, timeout_handler)  # 使用动态超时时间
+            timer = Timer(timeout, timeout_handler)
             timer.start()
 
             try:
@@ -291,7 +272,7 @@ def messages():
                     yield from send_event("content_block_delta", {
                         "type": "content_block_delta",
                         "index": 0,
-                        "delta": {"type": "text_delta", "text": "Request timed out. Partial response: " + ''.join(response_text)},
+                        "delta": {"type": "text_delta", "text": "Request timed out"},
                     })
                 
             except Exception as e:
@@ -306,7 +287,6 @@ 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",
@@ -364,11 +344,16 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout):
 
         sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
         
-        # Wait for response with timeout
         response_event.wait(timeout=timeout)
 
         if not response_text:
-            logger.warning("No response received (non-stream)", extra={'event_type': 'no_response_non_stream'})
+            logger.warning("No response received (non-stream)", extra={
+                'event_type': 'no_response_non_stream',
+                'data': {
+                    'input_tokens': input_tokens,
+                    'timeout': timeout
+                }
+            })
             return jsonify({"error": "No response received"}), 504
 
         full_response = {