Update app.py

app.py

@@ -8,6 +8,7 @@ import requests
 import logging
 from threading import Event, Timer
 import re
+import asyncio
 
 app = Flask(__name__)
 
@@ -98,9 +99,9 @@ def normalize_content(content):
     elif isinstance(content, dict):
         return json.dumps(content, ensure_ascii=False)
     elif isinstance(content, list):
-        return json.dumps(content, ensure_ascii=False)
+        return " ".join([normalize_content(item) for item in content])
     else:
-        return str(content)
+        return ""
 
 def calculate_tokens(text):
     if re.search(r'[^\x00-\x7F]', text):
@@ -136,6 +137,11 @@ def root():
         }
     })
 
+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()
@@ -144,21 +150,14 @@ def messages():
 
     try:
         json_body = request.json
-        model = json_body.get('model', 'claude-3-5-sonnet-20240620')
-        stream = json_body.get('stream', False)
-        messages = json_body.get('messages', [])
-
-        # 规范化所有消息
-        normalized_messages = [
-            {**msg, 'content': normalize_content(msg.get('content', ''))}
-            for msg in messages
-        ]
+        model = json_body.get('model', 'claude-3-opus-20240229')
+        stream = json_body.get('stream', True)
 
-        # 计算总输入 tokens
-        total_input_tokens = sum(calculate_tokens(msg['content']) for msg in normalized_messages)
+        previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
+        input_tokens = calculate_tokens(previous_messages)
 
-        # 准备完整的输入文本
-        full_input = "\n\n".join([msg['content'] for msg in normalized_messages])
+        # 根据 input_tokens 动态调整超时时间
+        timeout = max(30, min(300, input_tokens // 1000 * 30))  # 最少30秒，最多300秒
 
         msg_id = str(uuid.uuid4())
         response_event = Event()
@@ -166,21 +165,32 @@ def messages():
         response_text = []
         total_output_tokens = 0
 
-        # 动态调整超时时间
-        timeout_seconds = max(60, min(600, total_input_tokens // 1000 * 30))  # 每1000 tokens至少60秒，最多10分钟
+        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
 
-        def send_event(event_type, data):
-            event = create_event(event_type, data)
-            logger.info(f"Sending {event_type} event", extra={
-                'event_type': event_type,
-                'data': {'content': event}
-            })
-            return event
+        if not stream:
+            return handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout)
+
+        log_request(request.remote_addr, request.path, 200)
 
         def generate():
             nonlocal total_output_tokens
 
-            yield send_event("message_start", {
+            def send_event(event_type, data):
+                event = create_event(event_type, data)
+                logger.info(f"Sending {event_type} event", extra={
+                    'event_type': event_type,
+                    'data': {'content': event}
+                })
+                yield event
+
+            # Send initial events
+            yield from send_event("message_start", {
                 "type": "message_start",
                 "message": {
                     "id": msg_id,
@@ -190,127 +200,202 @@ def messages():
                     "content": [],
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": total_input_tokens, "output_tokens": 0},
+                    "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
                 },
             })
-            yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+            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"})
 
             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'])
-                    new_chunk = text['chunks'][-1] if text['chunks'] else None
-                    if new_chunk:
-                        response_text.append(new_chunk)
-                        chunk_tokens = calculate_tokens(new_chunk)
+                    chunk = text['chunks'][-1] if text['chunks'] else None
+                    if chunk:
+                        response_text.append(chunk)
+                        chunk_tokens = calculate_tokens(chunk)
                         total_output_tokens += chunk_tokens
-                        yield send_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": new_chunk},
+                        logger.info("Received chunk", extra={
+                            'event_type': 'chunk_received',
+                            'data': {
+                                'chunk': chunk,
+                                'tokens': chunk_tokens,
+                                '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={
+                        'event_type': 'response_complete',
+                        'data': {
+                            'total_tokens': total_output_tokens
+                        }
+                    })
                     response_event.set()
 
             def on_connect():
                 logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
-                sio.emit('perplexity_ask', (full_input, get_emit_data()))
-
-            def timeout_handler():
-                logger.warning(f"Request timed out after {timeout_seconds} seconds", extra={
-                    'event_type': 'request_timeout',
-                    'data': {'timeout_seconds': timeout_seconds}
+                emit_data = {
+                    "version": "2.9",
+                    "source": "default",
+                    "attachments": [],
+                    "language": "en-GB",
+                    "timezone": "Europe/London",
+                    "mode": "concise",
+                    "is_related_query": False,
+                    "is_default_related_query": False,
+                    "visitor_id": str(uuid.uuid4()),
+                    "frontend_context_uuid": str(uuid.uuid4()),
+                    "prompt_source": "user",
+                    "query_source": "home"
+                }
+                sio.emit('perplexity_ask', (previous_messages, emit_data))
+                logger.info("Sent query to Perplexity AI", extra={
+                    'event_type': 'query_sent',
+                    'data': {
+                        'message': previous_messages[:100] + '...' if len(previous_messages) > 100 else previous_messages
+                    }
                 })
-                timeout_event.set()
-                response_event.set()
 
             sio.on('connect', on_connect)
             sio.on('query_progress', on_query_progress)
 
-            timer = Timer(timeout_seconds, timeout_handler)
+            def timeout_handler():
+                logger.warning("Request timed out", extra={'event_type': 'request_timeout'})
+                timeout_event.set()
+                response_event.set()
+
+            timer = Timer(timeout, timeout_handler)  # 使用动态超时时间
             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():
                     sio.sleep(0.1)
-                    yield from on_query_progress({'text': json.dumps({'chunks': response_text})})
-                    response_text.clear()
-
+                    while response_text:
+                        chunk = response_text.pop(0)
+                        yield from send_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": chunk},
+                        })
+                
                 if timeout_event.is_set():
-                    yield send_event("content_block_delta", {
+                    yield from send_event("content_block_delta", {
                         "type": "content_block_delta",
                         "index": 0,
-                        "delta": {"type": "text_delta", "text": f"Request timed out after {timeout_seconds} seconds"},
+                        "delta": {"type": "text_delta", "text": "Request timed out. Partial response: " + ''.join(response_text)},
                     })
-
+                
             except Exception as e:
-                logger.error(f"Error during processing: {str(e)}", exc_info=True)
-                yield send_event("content_block_delta", {
+                logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
+                yield from send_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
-                    "delta": {"type": "text_delta", "text": f"Error during processing: {str(e)}"},
+                    "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
                 })
             finally:
                 timer.cancel()
                 if sio.connected:
                     sio.disconnect()
 
-            yield send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            yield send_event("message_delta", {
+            # 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",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                "usage": {"input_tokens": total_input_tokens, "output_tokens": total_output_tokens},
-            })
-            yield send_event("message_stop", {"type": "message_stop"})
-
-        if stream:
-            return Response(generate(), content_type='text/event-stream')
-        else:
-            # 非流式处理
-            full_response = []
-            for event in generate():
-                if 'content_block_delta' in event:
-                    data = json.loads(event.split('data: ')[1])
-                    full_response.append(data['delta']['text'])
-            
-            return jsonify({
-                "content": [{"text": ''.join(full_response), "type": "text"}],
-                "id": msg_id,
-                "model": model,
-                "role": "assistant",
-                "stop_reason": "end_turn",
-                "stop_sequence": None,
-                "type": "message",
-                "usage": {
-                    "input_tokens": total_input_tokens,
-                    "output_tokens": total_output_tokens,
-                },
+                "usage": {"output_tokens": total_output_tokens},
             })
+            yield from send_event("message_stop", {"type": "message_stop"})
+
+        return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
         logger.error(f"Request error: {str(e)}", exc_info=True)
-        logger.error(f"Request body: {json.dumps(request.json, default=str)}")
-        log_request(request.remote_addr, request.path, 500)
-        return jsonify({"error": str(e)}), 500
+        log_request(request.remote_addr, request.path, 400)
+        return jsonify({"error": str(e)}), 400
 
-def get_emit_data():
-    return {
-        "version": "2.9",
-        "source": "default",
-        "attachments": [],
-        "language": "en-GB",
-        "timezone": "Europe/London",
-        "mode": "concise",
-        "is_related_query": False,
-        "is_default_related_query": False,
-        "visitor_id": str(uuid.uuid4()),
-        "frontend_context_uuid": str(uuid.uuid4()),
-        "prompt_source": "user",
-        "query_source": "home"
-    }
+def handle_non_stream(previous_messages, msg_id, model, input_tokens, timeout):
+    try:
+        response_event = Event()
+        response_text = []
+        total_output_tokens = 0
+
+        def on_query_progress(data):
+            nonlocal total_output_tokens, response_text
+            if 'text' in data:
+                text = json.loads(data['text'])
+                chunk = text['chunks'][-1] if text['chunks'] else None
+                if chunk:
+                    response_text.append(chunk)
+                    chunk_tokens = calculate_tokens(chunk)
+                    total_output_tokens += chunk_tokens
+
+            if data.get('final', False):
+                response_event.set()
+
+        def on_connect():
+            logger.info("Connected to Perplexity AI (non-stream)", extra={'event_type': 'connection_established_non_stream'})
+            emit_data = {
+                "version": "2.9",
+                "source": "default",
+                "attachments": [],
+                "language": "en-GB",
+                "timezone": "Europe/London",
+                "mode": "concise",
+                "is_related_query": False,
+                "is_default_related_query": False,
+                "visitor_id": str(uuid.uuid4()),
+                "frontend_context_uuid": str(uuid.uuid4()),
+                "prompt_source": "user",
+                "query_source": "home"
+            }
+            sio.emit('perplexity_ask', (previous_messages, emit_data))
+
+        sio.on('connect', on_connect)
+        sio.on('query_progress', on_query_progress)
+
+        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'})
+            return jsonify({"error": "No response received"}), 504
+
+        full_response = {
+            "content": [{"text": ''.join(response_text), "type": "text"}],
+            "id": msg_id,
+            "model": model,
+            "role": "assistant",
+            "stop_reason": "end_turn",
+            "stop_sequence": None,
+            "type": "message",
+            "usage": {
+                "input_tokens": input_tokens,
+                "output_tokens": total_output_tokens,
+            },
+        }
+        logger.info("Sending non-stream response", extra={
+            'event_type': 'non_stream_response',
+            'data': {'content': full_response}
+        })
+        return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
+
+    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:
+        if sio.connected:
+            sio.disconnect()
 
 @app.errorhandler(404)
 def not_found(error):