Update app.py

app.py

@@ -9,16 +9,27 @@ import logging
 from threading import Event
 import re
 
+# 创建 Flask 应用
 app = Flask(__name__)
-logging.basicConfig(level=logging.INFO)
 
-# 从环境变量中获取API密钥
+# 自定义日志格式
+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_KEY = os.environ.get('PPLX_KEY')
 
 # 代理设置
 proxy_url = os.environ.get('PROXY_URL')
-
-# 设置代理
 if proxy_url:
     proxies = {
         'http': proxy_url,
@@ -29,7 +40,7 @@ if proxy_url:
 else:
     transport = None
 
-sio = socketio.Client(http_session=transport, logger=True, engineio_logger=True)
+sio = socketio.Client(http_session=transport, logger=sio_logger, engineio_logger=engineio_logger)
 
 # 连接选项
 connect_opts = {
@@ -49,7 +60,7 @@ sio_opts = {
 
 def log_request(ip, route, status):
     timestamp = datetime.now().isoformat()
-    logging.info(f"{timestamp} - {ip} - {route} - {status}")
+    app_logger.info(f"{timestamp} - {ip} - {route} - {status}")
 
 def validate_api_key():
     api_key = request.headers.get('x-api-key')
@@ -66,13 +77,10 @@ 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):
@@ -81,12 +89,9 @@ def calculate_tokens(text):
     - 对于英文和有空格的文本，使用空格分词。
     - 对于中文等没有空格的文本，使用字符级分词。
     """
-    # 首先判断文本是否包含大量非 ASCII 字符（如中文）
     if re.search(r'[^\x00-\x7F]', text):
-        # 如果包含非 ASCII 字符，使用字符级分词
         return len(text)
     else:
-        # 否则使用空格分词
         tokens = text.split()
         return len(tokens)
 
@@ -120,13 +125,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())
@@ -134,10 +136,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():
@@ -148,17 +148,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},  # 动态 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": ""}})
             yield create_event("ping", {"type": "ping"})
 
             def on_connect():
-                logging.info("Connected to Perplexity AI")
+                app_logger.info("Connected to Perplexity AI")
                 emit_data = {
                     "version": "2.9",
                     "source": "default",
@@ -183,25 +183,20 @@ def messages():
                     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")
+                app_logger.info("Disconnected from Perplexity AI")
                 response_event.set()
 
             def on_connect_error(data):
-                logging.error(f"Connection error: {data}")
+                app_logger.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)
 
@@ -219,7 +214,7 @@ def messages():
                         })
                 
             except Exception as e:
-                logging.error(f"Error during socket connection: {str(e)}")
+                app_logger.error(f"Error during socket connection: {str(e)}")
                 yield create_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
@@ -229,34 +224,30 @@ 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')
 
     except Exception as e:
-        logging.error(f"Request error: {str(e)}")
+        app_logger.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():
-            logging.info("Connected to Perplexity AI")
+            app_logger.info("Connected to Perplexity AI")
             emit_data = {
                 "version": "2.9",
                 "source": "default",
@@ -281,16 +272,15 @@ 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():
-            logging.info("Disconnected from Perplexity AI")
+            app_logger.info("Disconnected from Perplexity AI")
             response_event.set()
 
         def on_connect_error(data):
-            logging.error(f"Connection error: {data}")
+            app_logger.error(f"Connection error: {data}")
             response_text.append(f"Error connecting to Perplexity AI: {data}")
             response_event.set()
 
@@ -301,30 +291,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,  # 动态 input_tokens
-                "output_tokens": output_tokens,  # 动态 output_tokens
+                "input_tokens": input_tokens,
+                "output_tokens": output_tokens,
             },
         }
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
     except Exception as e:
-        logging.error(f"Error during socket connection: {str(e)}")
+        app_logger.error(f"Error during socket connection: {str(e)}")
         return jsonify({"error": str(e)}), 500
     finally:
         if sio.connected:
@@ -337,18 +323,18 @@ def not_found(error):
 
 @app.errorhandler(500)
 def server_error(error):
-    logging.error(f"Server error: {str(error)}")
+    app_logger.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))
-    logging.info(f"Perplexity proxy listening on port {port}")
+    app_logger.info(f"Perplexity proxy listening on port {port}")
     if not API_KEY:
-        logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
+        app_logger.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
     app.run(host='0.0.0.0', port=port)