Update app.py

app.py

@@ -1,18 +1,18 @@
 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
 import re
-import asyncio
-from functools import partial
+import math
 
 app = Flask(__name__)
 
+# 自定义日志格式化器
 class CustomFormatter(logging.Formatter):
     def format(self, record):
         log_data = {
@@ -39,20 +39,31 @@ 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:
     transport = None
 
-sio = socketio.AsyncClient(http_session=transport, logger=False, engineio_logger=False)
+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'),
@@ -96,7 +107,8 @@ def calculate_tokens(text):
     if re.search(r'[^\x00-\x7F]', text):
         return len(text)
     else:
-        return len(text.split())
+        tokens = text.split()
+        return len(tokens)
 
 def create_event(event, data):
     if isinstance(data, dict):
@@ -118,13 +130,34 @@ 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)"
                 }
             }
         }
     })
 
+# 在文件开头添加这个函数
+def split_messages(messages, max_tokens_per_chunk=8000):
+    chunks = []
+    current_chunk = []
+    current_chunk_tokens = 0
+
+    for message in messages:
+        message_tokens = calculate_tokens(message['content'])
+        if current_chunk_tokens + message_tokens > max_tokens_per_chunk and current_chunk:
+            chunks.append(current_chunk)
+            current_chunk = []
+            current_chunk_tokens = 0
+        
+        current_chunk.append(message)
+        current_chunk_tokens += message_tokens
+
+    if current_chunk:
+        chunks.append(current_chunk)
+
+    return chunks
+
 @app.route('/ai/v1/messages', methods=['POST'])
 def messages():
     auth_error = validate_api_key()
@@ -135,34 +168,25 @@ def messages():
         json_body = request.json
         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)
+        messages = json_body.get('messages', [])
+
+        # 分块处理
+        chunks = split_messages(messages)
+        total_chunks = len(chunks)
+        logger.info(f"Input split into {total_chunks} chunks", extra={
+            'event_type': 'input_split',
+            'data': {'total_chunks': total_chunks}
+        })
 
         msg_id = str(uuid.uuid4())
-        
-        if not stream:
-            return handle_non_stream(previous_messages, msg_id, model, input_tokens)
-
-        log_request(request.remote_addr, request.path, 200)
-
-        async def run_socket_io():
-            response_event = asyncio.Event()
-            response_text = []
-            total_output_tokens = 0
-            start_time = time.time()
-            last_activity_time = start_time
-            timeout = max(300, input_tokens / 100)  # 动态设置超时时间，最少300秒
-
-            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
+        total_input_tokens = sum(calculate_tokens(msg['content']) for msg in messages)
+        total_output_tokens = 0
+        full_response = []
+
+        def generate():
+            nonlocal total_output_tokens
 
-            yield send_event("message_start", {
+            yield from send_event("message_start", {
                 "type": "message_start",
                 "message": {
                     "id": msg_id,
@@ -172,148 +196,142 @@ def messages():
                     "content": [],
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
+                    "usage": {"input_tokens": total_input_tokens, "output_tokens": 0},
                 },
             })
-            yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
-            yield send_event("ping", {"type": "ping"})
-
-            @sio.event
-            async def query_progress(data):
-                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
-                    if chunk:
-                        response_text.append(chunk)
-                        chunk_tokens = calculate_tokens(chunk)
-                        total_output_tokens += chunk_tokens
-                        logger.info("Received chunk", extra={
-                            'event_type': 'chunk_received',
-                            'data': {
-                                'chunk': chunk,
-                                'tokens': chunk_tokens,
-                                'total_tokens': total_output_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()
 
-            @sio.event
-            async def connect():
-                logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
-                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"
-                }
-                await 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
-                    }
+            for chunk_index, chunk in enumerate(chunks):
+                chunk_input = "\n\n".join([normalize_content(msg['content']) for msg in chunk])
+                chunk_input_tokens = calculate_tokens(chunk_input)
+                
+                response_event = Event()
+                timeout_event = Event()
+                response_text = []
+
+                # 动态调整超时时间
+                timeout_seconds = max(30, min(300, chunk_input_tokens // 1000 * 30))
+
+                yield from send_event("chunk_start", {
+                    "type": "chunk_start",
+                    "chunk_index": chunk_index,
+                    "total_chunks": total_chunks,
                 })
 
-            async def heartbeat():
-                while not response_event.is_set():
-                    await sio.emit('ping')
-                    await asyncio.sleep(25)
+                def on_query_progress(data):
+                    nonlocal 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)
 
-            try:
-                await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-                heartbeat_task = asyncio.create_task(heartbeat())
-                
-                while not response_event.is_set() and (time.time() - start_time) < timeout:
-                    current_time = time.time()
-                    if current_time - last_activity_time > 60:  # 如果60秒内没有活动，记录警告
-                        logger.warning("No activity for 60 seconds", extra={'event_type': 'inactivity_warning'})
-                    await asyncio.sleep(0.1)
-                    while response_text:
-                        chunk = response_text.pop(0)
-                        yield send_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
-
-                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
-                        }
+                    if data.get('final', False):
+                        response_event.set()
+
+                sio.on('query_progress', on_query_progress)
+
+                def timeout_handler():
+                    logger.warning(f"Chunk {chunk_index + 1}/{total_chunks} timed out after {timeout_seconds} seconds", extra={
+                        'event_type': 'chunk_timeout',
+                        'data': {'chunk_index': chunk_index, 'total_chunks': total_chunks, 'timeout_seconds': timeout_seconds}
                     })
-                    yield send_event("content_block_delta", {
+                    timeout_event.set()
+                    response_event.set()
+
+                timer = Timer(timeout_seconds, timeout_handler)
+                timer.start()
+
+                try:
+                    sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+                    sio.emit('perplexity_ask', (chunk_input, get_emit_data()))
+
+                    while not response_event.is_set() and not timeout_event.is_set():
+                        sio.sleep(0.1)
+                        while response_text:
+                            new_chunk = response_text.pop(0)
+                            full_response.append(new_chunk)
+                            chunk_tokens = calculate_tokens(new_chunk)
+                            total_output_tokens += chunk_tokens
+                            yield from send_event("content_block_delta", {
+                                "type": "content_block_delta",
+                                "index": 0,
+                                "delta": {"type": "text_delta", "text": new_chunk},
+                            })
+
+                except Exception as e:
+                    logger.error(f"Error during chunk {chunk_index + 1}/{total_chunks} processing: {str(e)}", exc_info=True)
+                    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"},
+                        "delta": {"type": "text_delta", "text": f"Error processing chunk {chunk_index + 1}/{total_chunks}: {str(e)}"},
                     })
-            
-            except Exception as e:
-                logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
-                yield send_event("content_block_delta", {
-                    "type": "content_block_delta",
-                    "index": 0,
-                    "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
+                finally:
+                    timer.cancel()
+                    if sio.connected:
+                        sio.disconnect()
+
+                yield from send_event("chunk_end", {
+                    "type": "chunk_end",
+                    "chunk_index": chunk_index,
+                    "total_chunks": total_chunks,
                 })
-            finally:
-                heartbeat_task.cancel()
-                if sio.connected:
-                    await sio.disconnect()
 
-            yield send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            yield send_event("message_delta", {
+            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": {"output_tokens": total_output_tokens},
+                "usage": {"input_tokens": total_input_tokens, "output_tokens": total_output_tokens},
+            })
+            yield from send_event("message_stop", {"type": "message_stop"})
+
+        if stream:
+            return Response(generate(), content_type='text/event-stream')
+        else:
+            # 非流式处理
+            for _ in generate():
+                pass  # 处理所有生成的事件，但不发送
+            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,
+                },
             })
-            yield send_event("message_stop", {"type": "message_stop"})
-
-        def run_async():
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            try:
-                return loop.run_until_complete(run_socket_io())
-            finally:
-                loop.close()
-
-        return Response(run_async(), content_type='text/event-stream')
 
     except Exception as e:
         logger.error(f"Request error: {str(e)}", exc_info=True)
         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):
-    async def run_non_stream():
+    try:
+        response_event = Event()
         response_text = []
         total_output_tokens = 0
-        start_time = time.time()
-        timeout = max(300, input_tokens / 100)  # 动态设置超时时间，最少300秒
 
-        @sio.event
-        async def query_progress(data):
+        def on_query_progress(data):
             nonlocal total_output_tokens, response_text
             if 'text' in data:
                 text = json.loads(data['text'])
@@ -322,20 +340,11 @@ 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):
-                return True
+                response_event.set()
 
-        @sio.event
-        async def connect():
+        def on_connect():
             logger.info("Connected to Perplexity AI (non-stream)", extra={'event_type': 'connection_established_non_stream'})
             emit_data = {
                 "version": "2.9",
@@ -351,34 +360,19 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
                 "prompt_source": "user",
                 "query_source": "home"
             }
-            await sio.emit('perplexity_ask', (previous_messages, emit_data))
-
-        try:
-            await sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-            await asyncio.wait_for(sio.wait(), timeout=timeout)
-        except asyncio.TimeoutError:
-            logger.warning(f"Request timed out after {timeout} seconds (non-stream)", extra={
-                'event_type': 'request_timeout_non_stream',
-                'data': {
-                    'timeout': timeout,
-                    'input_tokens': input_tokens,
-                    'elapsed_time': time.time() - start_time
-                }
-            })
-        finally:
-            if sio.connected:
-                await sio.disconnect()
+            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=30)
 
         if not response_text:
-            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
+            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"}],
@@ -395,22 +389,16 @@ 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,
-                'elapsed_time': time.time() - start_time
-            }
+            'data': {'content': full_response}
         })
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    try:
-        return loop.run_until_complete(run_non_stream())
     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:
-        loop.close()
+        if sio.connected:
+            sio.disconnect()
 
 @app.errorhandler(404)
 def not_found(error):