Update app.py

app.py

@@ -11,18 +11,11 @@ from threading import Event
 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:
@@ -30,12 +23,7 @@ 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'),
@@ -58,31 +46,24 @@ 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):
-    """
-    计算文本的 token 数量。
-    这里我们简单地通过空格分词来模拟 token 计数。
-    如果使用 GPT 模型，可以使用 tiktoken 库进行 tokenization。
-    """
-    # 使用简单的空格分词计数
-    tokens = text.split()
-    return len(tokens)
+    return len(text.split())
+
+def validate_json(data):
+    try:
+        json.loads(json.dumps(data))
+        return True
+    except json.JSONDecodeError:
+        return False
 
 @app.route('/')
 def root():
@@ -114,13 +95,10 @@ def messages():
 
     try:
         json_body = request.json
-        model = json_body.get('model', 'claude-3-opus-20240229')  # 动态获取模型，默认 claude-3-opus-20240229
-        stream = json_body.get('stream', True)  # 默认为True
+        model = json_body.get('model', 'claude-3-opus-20240229')
+        stream = json_body.get('stream', 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())
@@ -128,10 +106,8 @@ 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():
@@ -142,10 +118,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},  # 动态 input_tokens
+                    "usage": {"input_tokens": input_tokens, "output_tokens": 1},
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -171,23 +147,16 @@ def messages():
 
             def on_query_progress(data):
                 nonlocal response_text
-                if 'text' in data:
-                    try:
-                        # 解析嵌套的 JSON 字符串
-                        text_content = json.loads(data['text'])  # data['text'] 是一个 JSON 字符串
-                        chunk = text_content['chunks'][-1] if 'chunks' in text_content and text_content['chunks'] else None
+                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)
-                    except json.JSONDecodeError as e:
-                        logging.error(f"Failed to decode JSON from 'text' field: {e}")
-                        response_text.append(f"Error decoding response: {e}")
-
-                # 检查是否是最终响应
-                if data.get('final', False):
-                    response_event.set()
-
-            def on_query_complete(data):
-                response_event.set()
+                    if data.get('final', False):
+                        response_event.set()
+                except json.JSONDecodeError:
+                    logging.error(f"Failed to parse query progress data: {data}")
 
             def on_disconnect():
                 logging.info("Disconnected from Perplexity AI")
@@ -200,7 +169,6 @@ 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)
 
@@ -211,11 +179,15 @@ def messages():
                     sio.sleep(0.1)
                     while response_text:
                         chunk = response_text.pop(0)
-                        yield create_event("content_block_delta", {
+                        event_data = {
                             "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)}")
@@ -228,16 +200,15 @@ 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},  # 动态 output_tokens
+                "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
             })
-            yield create_event("message_stop", {"type": "message_stop"})  # 确保发送 message_stop 事件
+            yield create_event("message_stop", {"type": "message_stop"})
 
         return Response(generate(), content_type='text/event-stream')
 
@@ -247,9 +218,6 @@ 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 = []
@@ -274,20 +242,16 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
 
         def on_query_progress(data):
             nonlocal response_text
-            if 'text' in data:
-                try:
-                    # 解析嵌套的 JSON 字符串
-                    text_content = json.loads(data['text'])  # data['text'] 是一个 JSON 字符串
-                    chunk = text_content['chunks'][-1] if 'chunks' in text_content and text_content['chunks'] else None
+            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)
-                except json.JSONDecodeError as e:
-                    logging.error(f"Failed to decode JSON from 'text' field: {e}")
-                    response_text.append(f"Error decoding response: {e}")
-
-            # 检查是否是最终响应
-            if data.get('final', False):
-                response_event.set()
+                if data.get('final', False):
+                    response_event.set()
+            except json.JSONDecodeError:
+                logging.error(f"Failed to parse query progress data: {data}")
 
         def on_disconnect():
             logging.info("Disconnected from Perplexity AI")
@@ -305,26 +269,28 @@ 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,  # 动态 input_tokens
-                "output_tokens": output_tokens,  # 动态 output_tokens
+                "input_tokens": input_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:
@@ -346,9 +312,13 @@ def server_error(error):
     return "Something broke!", 500
 
 def create_event(event, data):
-    if isinstance(data, dict):
-        data = json.dumps(data, ensure_ascii=False)  # 确保中文不会被转义
-    return f"event: {event}\ndata: {data}\n\n"
+    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 __name__ == '__main__':
     port = int(os.environ.get('PORT', 8081))