Update app.py

app.py

@@ -8,6 +8,7 @@ import requests
 import logging
 from threading import Event
 import re
+from functools import wraps
 
 # 创建 Flask 应用
 app = Flask(__name__)
@@ -15,37 +16,21 @@ app = Flask(__name__)
 # 自定义日志格式
 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')
+transport = requests.Session()
 if proxy_url:
-    proxies = {
-        'http': proxy_url,
-        'https': proxy_url
-    }
-    transport = requests.Session()
-    transport.proxies.update(proxies)
-else:
-    transport = None
+    transport.proxies.update({'http': proxy_url, 'https': proxy_url})
 
-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 = {
-    'transports': ['websocket', 'polling'],  # 允许回退到轮询
-}
+connect_opts = {'transports': ['websocket', 'polling']}
 
 # 其他选项
 sio_opts = {
@@ -58,52 +43,45 @@ sio_opts = {
     }
 }
 
-def log_request(ip, route, status):
-    timestamp = datetime.now().isoformat()
-    app_logger.info(f"{timestamp} - {ip} - {route} - {status}")
-
-def validate_api_key():
-    api_key = request.headers.get('x-api-key')
-    if api_key != API_KEY:
-        log_request(request.remote_addr, request.path, 401)
-        return create_json_response({"error": "Invalid API key"}), 401
-    return None
+def log_request(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        start_time = datetime.now()
+        response = func(*args, **kwargs)
+        duration = (datetime.now() - start_time).total_seconds()
+        app_logger.info(f"{request.remote_addr} - {request.method} {request.path} - {response.status_code} - {duration:.2f}s")
+        return response
+    return wrapper
+
+def validate_api_key(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        api_key = request.headers.get('x-api-key')
+        if api_key != API_KEY:
+            app_logger.warning(f"Invalid API key attempt from {request.remote_addr}")
+            return jsonify({"error": "Invalid API key"}), 401
+        return func(*args, **kwargs)
+    return wrapper
 
 def normalize_content(content):
-    """
-    递归处理 msg['content']，确保其为字符串。
-    如果 content 是字典或列表，将其转换为字符串。
-    """
     if isinstance(content, str):
         return content
-    elif isinstance(content, dict):
+    elif isinstance(content, (dict, list)):
         return json.dumps(content, ensure_ascii=False)
-    elif isinstance(content, list):
-        return " ".join([normalize_content(item) for item in content])
-    else:
-        return ""
+    return str(content)
 
 def calculate_tokens(text):
-    """
-    改进的 token 计算方法。
-    - 对于英文和有空格的文本，使用空格分词。
-    - 对于中文等没有空格的文本，使用字符级分词。
-    """
-    if re.search(r'[^\x00-\x7F]', text):
-        return len(text)
-    else:
-        tokens = text.split()
-        return len(tokens)
-
-def create_json_response(data):
-    """创建一个JSON响应，并在发送之前记录它"""
-    json_str = json.dumps(data, ensure_ascii=False)
-    app_logger.info(f"Sending JSON response: {json_str}")
-    return Response(json_str, content_type='application/json')
+    return len(re.findall(r'\w+|[^\w\s]', text, re.UNICODE))
+
+def create_json_response(data, status_code=200):
+    response = jsonify(data)
+    response.status_code = status_code
+    app_logger.debug(f"Sending JSON response: {json.dumps(data, ensure_ascii=False)}")
+    return response
 
 @app.route('/')
+@log_request
 def root():
-    log_request(request.remote_addr, request.path, 200)
     return create_json_response({
         "message": "Welcome to the Perplexity AI Proxy API",
         "endpoints": {
@@ -124,11 +102,9 @@ def root():
     })
 
 @app.route('/ai/v1/messages', methods=['POST'])
+@log_request
+@validate_api_key
 def messages():
-    auth_error = validate_api_key()
-    if auth_error:
-        return auth_error
-
     try:
         json_body = request.json
         model = json_body.get('model', 'claude-3-opus-20240229')
@@ -144,11 +120,9 @@ def messages():
         if not stream:
             return handle_non_stream(previous_messages, msg_id, model, input_tokens)
 
-        log_request(request.remote_addr, request.path, 200)
-
         def generate():
             try:
-                yield create_event("message_start", {
+                yield create_sse_event("message_start", {
                     "type": "message_start",
                     "message": {
                         "id": msg_id,
@@ -161,9 +135,12 @@ def messages():
                         "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"})
+                yield create_sse_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+                yield create_sse_event("ping", {"type": "ping"})
 
+                sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+                
+                @sio.on('connect')
                 def on_connect():
                     app_logger.info("Connected to Perplexity AI")
                     emit_data = {
@@ -182,73 +159,67 @@ def messages():
                     }
                     sio.emit('perplexity_ask', (previous_messages, emit_data))
 
+                @sio.on('query_progress')
                 def on_query_progress(data):
-                    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)
+                            yield create_sse_event("content_block_delta", {
+                                "type": "content_block_delta",
+                                "index": 0,
+                                "delta": {"type": "text_delta", "text": chunk},
+                            })
 
                     if data.get('final', False):
                         response_event.set()
 
+                @sio.on('disconnect')
                 def on_disconnect():
                     app_logger.info("Disconnected from Perplexity AI")
                     response_event.set()
 
+                @sio.on('connect_error')
                 def on_connect_error(data):
                     app_logger.error(f"Connection error: {data}")
-                    response_text.append(f"Error connecting to Perplexity AI: {data}")
+                    yield create_sse_event("error", {"type": "error", "message": f"Error connecting to Perplexity AI: {data}"})
                     response_event.set()
 
-                sio.on('connect', on_connect)
-                sio.on('query_progress', on_query_progress)
-                sio.on('disconnect', on_disconnect)
-                sio.on('connect_error', on_connect_error)
-
-                sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-                
                 while not response_event.is_set():
                     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,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
-                
+
                 output_tokens = calculate_tokens(''.join(response_text))
 
-                yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-                yield create_event("message_delta", {
+                yield create_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+                yield create_sse_event("message_delta", {
                     "type": "message_delta",
                     "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                     "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
                 })
-                yield create_event("message_stop", {"type": "message_stop"})
+                yield create_sse_event("message_stop", {"type": "message_stop"})
 
             except Exception as e:
                 app_logger.error(f"Error in generate function: {str(e)}")
-                yield create_event("error", {"type": "error", "message": str(e)})
+                yield create_sse_event("error", {"type": "error", "message": str(e)})
             finally:
                 if sio.connected:
                     sio.disconnect()
 
-        return Response(generate(), content_type='text/event-stream')
+        return Response(generate(), content_type='text/event-stream; charset=utf-8')
 
     except Exception as e:
         app_logger.error(f"Request error: {str(e)}")
-        log_request(request.remote_addr, request.path, 400)
-        return create_json_response({"error": str(e)}), 400
+        return create_json_response({"error": str(e)}, 400)
 
 def handle_non_stream(previous_messages, msg_id, model, input_tokens):
     try:
         response_event = Event()
         response_text = []
 
+        sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+
+        @sio.on('connect')
         def on_connect():
             app_logger.info("Connected to Perplexity AI")
             emit_data = {
@@ -267,8 +238,8 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
             }
             sio.emit('perplexity_ask', (previous_messages, emit_data))
 
+        @sio.on('query_progress')
         def on_query_progress(data):
-            nonlocal response_text
             if 'text' in data:
                 text = json.loads(data['text'])
                 chunk = text['chunks'][-1] if text['chunks'] else None
@@ -278,22 +249,17 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
             if data.get('final', False):
                 response_event.set()
 
+        @sio.on('disconnect')
         def on_disconnect():
             app_logger.info("Disconnected from Perplexity AI")
             response_event.set()
 
+        @sio.on('connect_error')
         def on_connect_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('disconnect', on_disconnect)
-        sio.on('connect_error', on_connect_error)
-
-        sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
-        
         response_event.wait(timeout=30)
         output_tokens = calculate_tokens(''.join(response_text))
 
@@ -315,27 +281,24 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
 
     except Exception as e:
         app_logger.error(f"Error during socket connection: {str(e)}")
-        return create_json_response({"error": str(e)}), 500
+        return create_json_response({"error": str(e)}, 500)
     finally:
         if sio.connected:
             sio.disconnect()
 
 @app.errorhandler(404)
 def not_found(error):
-    log_request(request.remote_addr, request.path, 404)
-    return create_json_response({"error": "Not Found"}), 404
+    return create_json_response({"error": "Not Found"}, 404)
 
 @app.errorhandler(500)
 def server_error(error):
     app_logger.error(f"Server error: {str(error)}")
-    log_request(request.remote_addr, request.path, 500)
-    return create_json_response({"error": "Internal Server Error"}), 500
-
-def create_event(event, data):
-    if isinstance(data, dict):
-        data = json.dumps(data, ensure_ascii=False)
-    event_str = f"event: {event}\ndata: {data}\n\n"
-    app_logger.info(f"Sending SSE event: {event_str}")
+    return create_json_response({"error": "Internal Server Error"}, 500)
+
+def create_sse_event(event, data):
+    json_data = json.dumps(data, ensure_ascii=False)
+    event_str = f"event: {event}\ndata: {json_data}\n\n"
+    app_logger.debug(f"Sending SSE event: {event_str}")
     return event_str
 
 if __name__ == '__main__':
@@ -343,4 +306,4 @@ if __name__ == '__main__':
     app_logger.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.")
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port, debug=False)