Update app.py

app.py

@@ -8,14 +8,24 @@ import requests
 import logging
 from threading import Event
 
+# 如果使用 GPT 模型的 tokenization，可以引入 tiktoken
+# import tiktoken  # 如果需要使用 GPT 的 token 化库
+
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
 
+# 从环境变量中获取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:
@@ -23,7 +33,12 @@ else:
 
 sio = socketio.Client(http_session=transport, logger=True, engineio_logger=True)
 
-connect_opts = {'transports': ['websocket', 'polling']}
+# 连接选项
+connect_opts = {
+    'transports': ['websocket', 'polling'],  # 允许回退到轮询
+}
+
+# 其他选项
 sio_opts = {
     'extraHeaders': {
         'Cookie': os.environ.get('PPLX_COOKIE'),
@@ -46,24 +61,31 @@ def validate_api_key():
     return None
 
 def normalize_content(content):
+    """
+    递归处理 msg['content']，确保其为字符串。
+    如果 content 是字典或列表，将其转换为字符串。
+    """
     if isinstance(content, str):
         return content
     elif isinstance(content, dict):
+        # 将字典转化为 JSON 字符串
         return json.dumps(content, ensure_ascii=False)
     elif isinstance(content, list):
+        # 对于列表，递归处理每个元素
         return " ".join([normalize_content(item) for item in content])
     else:
+        # 如果是其他类型，返回空字符串
         return ""
 
 def calculate_tokens(text):
-    return len(text.split())
-
-def validate_json(data):
-    try:
-        json.loads(json.dumps(data))
-        return True
-    except json.JSONDecodeError:
-        return False
+    """
+    计算文本的 token 数量。
+    这里我们简单地通过空格分词来模拟 token 计数。
+    如果使用 GPT 模型，可以使用 tiktoken 库进行 tokenization。
+    """
+    # 使用简单的空格分词计数
+    tokens = text.split()
+    return len(tokens)
 
 @app.route('/')
 def root():
@@ -95,10 +117,13 @@ 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-opus-20240229')  # 动态获取模型，默认 claude-3-opus-20240229
+        stream = json_body.get('stream', True)  # 默认为True
 
+        # 使用 normalize_content 递归处理 msg['content']
         previous_messages = "\n\n".join([normalize_content(msg['content']) for msg in json_body['messages']])
+
+        # 动态计算输入的 token 数量
         input_tokens = calculate_tokens(previous_messages)
 
         msg_id = str(uuid.uuid4())
@@ -106,8 +131,10 @@ def messages():
         response_text = []
 
         if not stream:
+            # 处理 stream 为 false 的情况
             return handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
+        # 记录日志：此时请求上下文仍然有效
         log_request(request.remote_addr, request.path, 200)
 
         def generate():
@@ -118,10 +145,10 @@ def messages():
                     "type": "message",
                     "role": "assistant",
                     "content": [],
-                    "model": model,
+                    "model": model,  # 动态模型
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": 1},
+                    "usage": {"input_tokens": input_tokens, "output_tokens": 1},  # 动态 input_tokens
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -147,16 +174,18 @@ def messages():
 
             def on_query_progress(data):
                 nonlocal response_text
-                try:
-                    if 'text' in data:
-                        text = json.loads(data['text'])
-                        chunk = text['chunks'][-1] if text['chunks'] else None
-                        if chunk:
-                            response_text.append(chunk)
-                    if data.get('final', False):
-                        response_event.set()
-                except json.JSONDecodeError:
-                    logging.error(f"Failed to parse query progress data: {data}")
+                if 'text' in data:
+                    text = json.loads(data['text'])
+                    chunk = text['chunks'][-1] if text['chunks'] else None
+                    if chunk:
+                        response_text.append(chunk)
+
+                # 检查是否是最终响应
+                if data.get('final', False):
+                    response_event.set()
+
+            def on_query_complete(data):
+                response_event.set()
 
             def on_disconnect():
                 logging.info("Disconnected from Perplexity AI")
@@ -169,6 +198,7 @@ def messages():
 
             sio.on('connect', on_connect)
             sio.on('query_progress', on_query_progress)
+            sio.on('query_complete', on_query_complete)
             sio.on('disconnect', on_disconnect)
             sio.on('connect_error', on_connect_error)
 
@@ -179,15 +209,11 @@ def messages():
                     sio.sleep(0.1)
                     while response_text:
                         chunk = response_text.pop(0)
-                        event_data = {
+                        yield create_event("content_block_delta", {
                             "type": "content_block_delta",
                             "index": 0,
                             "delta": {"type": "text_delta", "text": chunk},
-                        }
-                        if validate_json(event_data):
-                            yield create_event("content_block_delta", event_data)
-                        else:
-                            logging.error(f"Invalid JSON for content_block_delta: {event_data}")
+                        })
                 
             except Exception as e:
                 logging.error(f"Error during socket connection: {str(e)}")
@@ -200,15 +226,16 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
+            # 动态计算输出的 token 数量
             output_tokens = calculate_tokens(''.join(response_text))
 
             yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
             yield create_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
+                "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},  # 动态 output_tokens
             })
-            yield create_event("message_stop", {"type": "message_stop"})
+            yield create_event("message_stop", {"type": "message_stop"})  # 确保发送 message_stop 事件
 
         return Response(generate(), content_type='text/event-stream')
 
@@ -218,6 +245,9 @@ def messages():
         return jsonify({"error": str(e)}), 400
 
 def handle_non_stream(previous_messages, msg_id, model, input_tokens):
+    """
+    处理 stream 为 false 的情况，返回完整的响应。
+    """
     try:
         response_event = Event()
         response_text = []
@@ -242,16 +272,15 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
 
         def on_query_progress(data):
             nonlocal response_text
-            try:
-                if 'text' in data:
-                    text = json.loads(data['text'])
-                    chunk = text['chunks'][-1] if text['chunks'] else None
-                    if chunk:
-                        response_text.append(chunk)
-                if data.get('final', False):
-                    response_event.set()
-            except json.JSONDecodeError:
-                logging.error(f"Failed to parse query progress data: {data}")
+            if 'text' in data:
+                text = json.loads(data['text'])
+                chunk = text['chunks'][-1] if text['chunks'] else None
+                if chunk:
+                    response_text.append(chunk)
+
+            # 检查是否是最终响应
+            if data.get('final', False):
+                response_event.set()
 
         def on_disconnect():
             logging.info("Disconnected from Perplexity AI")
@@ -269,28 +298,26 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
 
         sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
         
+        # 等待响应完成
         response_event.wait(timeout=30)
 
+        # 动态计算输出的 token 数量
         output_tokens = calculate_tokens(''.join(response_text))
 
+        # 生成完整的响应
         full_response = {
-            "content": [{"text": ''.join(response_text), "type": "text"}],
+            "content": [{"text": ''.join(response_text), "type": "text"}],  # 合并所有文本块
             "id": msg_id,
-            "model": model,
+            "model": model,  # 动态模型
             "role": "assistant",
             "stop_reason": "end_turn",
             "stop_sequence": None,
             "type": "message",
             "usage": {
-                "input_tokens": input_tokens,
-                "output_tokens": output_tokens,
+                "input_tokens": input_tokens,  # 动态 input_tokens
+                "output_tokens": output_tokens,  # 动态 output_tokens
             },
         }
-
-        if not validate_json(full_response):
-            logging.error(f"Invalid JSON response: {full_response}")
-            return jsonify({"error": "Invalid response format"}), 500
-
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
     except Exception as e:
@@ -312,13 +339,9 @@ def server_error(error):
     return "Something broke!", 500
 
 def create_event(event, data):
-    try:
-        if isinstance(data, dict):
-            data = json.dumps(data, ensure_ascii=False)
-        return f"event: {event}\ndata: {data}\n\n"
-    except json.JSONDecodeError:
-        logging.error(f"Failed to serialize event data: {data}")
-        return f"event: {event}\ndata: {json.dumps({'error': 'Data serialization failed'})}\n\n"
+    if isinstance(data, dict):
+        data = json.dumps(data, ensure_ascii=False)  # 确保中文不会被转义
+    return f"event: {event}\ndata: {data}\n\n"
 
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 8081))