Update app.py

app.py

@@ -9,27 +9,16 @@ import logging
 from threading import Event
 import re
 
-# 创建 Flask 应用
 app = Flask(__name__)
+logging.basicConfig(level=logging.INFO)
 
-# 自定义日志格式
-log_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
-logging.basicConfig(level=logging.INFO, format=log_format)
-
-# 创建不同的日志记录器
-app_logger = logging.getLogger('app')
-sio_logger = logging.getLogger('socketio.client')
-engineio_logger = logging.getLogger('engineio.client')
-
-# 调整日志级别，隐藏过多的调试信息
-sio_logger.setLevel(logging.WARNING)
-engineio_logger.setLevel(logging.WARNING)
-
-# 从环境变量中获取 API 密钥
+# 从环境变量中获取API密钥
 API_KEY = os.environ.get('PPLX_KEY')
 
 # 代理设置
 proxy_url = os.environ.get('PROXY_URL')
+
+# 设置代理
 if proxy_url:
     proxies = {
         'http': proxy_url,
@@ -40,7 +29,7 @@ if proxy_url:
 else:
     transport = None
 
-sio = socketio.Client(http_session=transport, logger=sio_logger, engineio_logger=engineio_logger)
+sio = socketio.Client(http_session=transport, logger=True, engineio_logger=True)
 
 # 连接选项
 connect_opts = {
@@ -60,7 +49,7 @@ sio_opts = {
 
 def log_request(ip, route, status):
     timestamp = datetime.now().isoformat()
-    app_logger.info(f"{timestamp} - {ip} - {route} - {status}")
+    logging.info(f"{timestamp} - {ip} - {route} - {status}")
 
 def validate_api_key():
     api_key = request.headers.get('x-api-key')
@@ -77,10 +66,13 @@ def normalize_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):
@@ -89,9 +81,12 @@ def calculate_tokens(text):
     - 对于英文和有空格的文本，使用空格分词。
     - 对于中文等没有空格的文本，使用字符级分词。
     """
+    # 首先判断文本是否包含大量非 ASCII 字符（如中文）
     if re.search(r'[^\x00-\x7F]', text):
+        # 如果包含非 ASCII 字符，使用字符级分词
         return len(text)
     else:
+        # 否则使用空格分词
         tokens = text.split()
         return len(tokens)
 
@@ -125,19 +120,25 @@ 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())
         response_event = Event()
         response_text = []
+        output_tokens = 0  # 初始化 output_tokens
 
         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():
@@ -148,17 +149,17 @@ 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": 0},  # 动态 input_tokens，初始化 output_tokens
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
             yield create_event("ping", {"type": "ping"})
 
             def on_connect():
-                app_logger.info("Connected to Perplexity AI")
+                logging.info("Connected to Perplexity AI")
                 emit_data = {
                     "version": "2.9",
                     "source": "default",
@@ -176,27 +177,41 @@ def messages():
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
 
             def on_query_progress(data):
-                nonlocal response_text
+                nonlocal response_text, output_tokens
                 if 'text' in data:
                     text = json.loads(data['text'])
                     chunk = text['chunks'][-1] if text['chunks'] else None
                     if chunk:
                         response_text.append(chunk)
+                        # 动态计算当前块的 token 数量并累积到 output_tokens
+                        output_tokens += calculate_tokens(chunk)
+
+                        # 发送 message_delta 事件，包含当前累积的 output_tokens
+                        yield create_event("message_delta", {
+                            "type": "message_delta",
+                            "delta": {"stop_reason": None, "stop_sequence": None},
+                            "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},  # 累积 output_tokens
+                        })
 
+                # 检查是否是最终响应
                 if data.get('final', False):
                     response_event.set()
 
+            def on_query_complete(data):
+                response_event.set()
+
             def on_disconnect():
-                app_logger.info("Disconnected from Perplexity AI")
+                logging.info("Disconnected from Perplexity AI")
                 response_event.set()
 
             def on_connect_error(data):
-                app_logger.error(f"Connection error: {data}")
+                logging.error(f"Connection error: {data}")
                 response_text.append(f"Error connecting to Perplexity AI: {data}")
                 response_event.set()
 
             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)
 
@@ -207,8 +222,6 @@ def messages():
                     sio.sleep(0.1)
                     while response_text:
                         chunk = response_text.pop(0)
-                        # 记录每一块返回的内容
-                        app_logger.info(f"Returning chunk to client: {chunk}")
                         yield create_event("content_block_delta", {
                             "type": "content_block_delta",
                             "index": 0,
@@ -216,7 +229,7 @@ def messages():
                         })
                 
             except Exception as e:
-                app_logger.error(f"Error during socket connection: {str(e)}")
+                logging.error(f"Error during socket connection: {str(e)}")
                 yield create_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
@@ -226,30 +239,32 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            output_tokens = calculate_tokens(''.join(response_text))
-
+            # 发送最终的 message_stop 事件，包含最终的 output_tokens
             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')
 
     except Exception as e:
-        app_logger.error(f"Request error: {str(e)}")
+        logging.error(f"Request error: {str(e)}")
         log_request(request.remote_addr, request.path, 400)
         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 = []
 
         def on_connect():
-            app_logger.info("Connected to Perplexity AI")
+            logging.info("Connected to Perplexity AI")
             emit_data = {
                 "version": "2.9",
                 "source": "default",
@@ -274,15 +289,16 @@ 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()
 
         def on_disconnect():
-            app_logger.info("Disconnected from Perplexity AI")
+            logging.info("Disconnected from Perplexity AI")
             response_event.set()
 
         def on_connect_error(data):
-            app_logger.error(f"Connection error: {data}")
+            logging.error(f"Connection error: {data}")
             response_text.append(f"Error connecting to Perplexity AI: {data}")
             response_event.set()
 
@@ -293,29 +309,30 @@ 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
             },
         }
-        
-        # 记录完整的响应内容
-        app_logger.info(f"Returning full response to client: {full_response}")
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
     except Exception as e:
-        app_logger.error(f"Error during socket connection: {str(e)}")
+        logging.error(f"Error during socket connection: {str(e)}")
         return jsonify({"error": str(e)}), 500
     finally:
         if sio.connected:
@@ -328,18 +345,18 @@ def not_found(error):
 
 @app.errorhandler(500)
 def server_error(error):
-    app_logger.error(f"Server error: {str(error)}")
+    logging.error(f"Server error: {str(error)}")
     log_request(request.remote_addr, request.path, 500)
     return "Something broke!", 500
 
 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))
-    app_logger.info(f"Perplexity proxy listening on port {port}")
+    logging.info(f"Perplexity proxy listening on port {port}")
     if not API_KEY:
-        app_logger.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
+        logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
     app.run(host='0.0.0.0', port=port)