Update app.py

app.py

@@ -59,19 +59,34 @@ 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 计算方法。
+    - 对于英文和有空格的文本，使用空格分词。
+    - 对于中文等没有空格的文本，使用字符级分词。
+    """
+    # 首先判断文本是否包含大量非 ASCII 字符（如中文）
     if re.search(r'[^\x00-\x7F]', text):
+        # 如果包含非 ASCII 字符，使用字符级分词
         return len(text)
     else:
+        # 否则使用空格分词
         tokens = text.split()
         return len(tokens)
 
@@ -105,10 +120,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())
@@ -116,13 +134,13 @@ 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():
-            output_tokens = 0  # 初始化 output_tokens
-
             yield create_event("message_start", {
                 "type": "message_start",
                 "message": {
@@ -130,10 +148,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": output_tokens},
+                    "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": ""}})
@@ -158,24 +176,14 @@ def messages():
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
 
             def on_query_progress(data):
-                nonlocal response_text, output_tokens
+                nonlocal 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)
-                        output_tokens += chunk_tokens
-                        yield create_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
-                        yield create_event("message_delta", {
-                            "type": "message_delta",
-                            "delta": {"usage": {"output_tokens": output_tokens}},
-                        })
 
+                # 检查是否是最终响应
                 if data.get('final', False):
                     response_event.set()
 
@@ -204,7 +212,11 @@ def messages():
                     sio.sleep(0.1)
                     while response_text:
                         chunk = response_text.pop(0)
-                        yield from on_query_progress({"text": json.dumps({"chunks": [chunk]})})
+                        yield create_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": chunk},
+                        })
                 
             except Exception as e:
                 logging.error(f"Error during socket connection: {str(e)}")
@@ -217,13 +229,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')
 
@@ -233,6 +248,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 = []
@@ -263,6 +281,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
                 if chunk:
                     response_text.append(chunk)
 
+            # 检查是否是最终响应
             if data.get('final', False):
                 response_event.set()
 
@@ -282,21 +301,24 @@ 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
             },
         }
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
@@ -321,10 +343,9 @@ def server_error(error):
 
 def create_event(event, data):
     if isinstance(data, dict):
-        data = json.dumps(data, ensure_ascii=False)
+        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))
     logging.info(f"Perplexity proxy listening on port {port}")