Update app.py

app.py

@@ -9,8 +9,29 @@ import logging
 from threading import Event
 import re
 
+# 自定义日志格式化器
+class CustomFormatter(logging.Formatter):
+    def format(self, record):
+        log_data = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "message": record.getMessage(),
+        }
+        if hasattr(record, 'event_type'):
+            log_data['event_type'] = record.event_type
+        if hasattr(record, 'data'):
+            log_data['data'] = record.data
+        return json.dumps(log_data, ensure_ascii=False, indent=2)
+
+def setup_logging():
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler()
+    handler.setFormatter(CustomFormatter())
+    logger.addHandler(handler)
+
 app = Flask(__name__)
-logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # 从环境变量中获取API密钥
 API_KEY = os.environ.get('PPLX_KEY')
@@ -29,7 +50,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=False, engineio_logger=False)
 
 # 连接选项
 connect_opts = {
@@ -48,8 +69,14 @@ sio_opts = {
 }
 
 def log_request(ip, route, status):
-    timestamp = datetime.now().isoformat()
-    logging.info(f"{timestamp} - {ip} - {route} - {status}")
+    logger.info(f"HTTP Request", extra={
+        'event_type': 'http_request',
+        'data': {
+            'ip': ip,
+            'route': route,
+            'status': status
+        }
+    })
 
 def validate_api_key():
     api_key = request.headers.get('x-api-key')
@@ -154,12 +181,26 @@ def messages():
                         response_text.append(chunk)
                         chunk_tokens = calculate_tokens(chunk)
                         total_output_tokens += chunk_tokens
+                        logger.info("Received chunk", extra={
+                            'event_type': 'chunk_received',
+                            'data': {
+                                'chunk': chunk,
+                                'tokens': chunk_tokens,
+                                'total_tokens': total_output_tokens
+                            }
+                        })
 
                 if data.get('final', False):
+                    logger.info("Final response received", extra={
+                        'event_type': 'response_complete',
+                        'data': {
+                            'total_tokens': total_output_tokens
+                        }
+                    })
                     response_event.set()
 
             def on_connect():
-                logging.info("Connected to Perplexity AI")
+                logger.info("Connected to Perplexity AI", extra={'event_type': 'connection_established'})
                 emit_data = {
                     "version": "2.9",
                     "source": "default",
@@ -175,6 +216,12 @@ def messages():
                     "query_source": "home"
                 }
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
+                logger.info("Sent query to Perplexity AI", extra={
+                    'event_type': 'query_sent',
+                    'data': {
+                        'message': previous_messages[:100] + '...' if len(previous_messages) > 100 else previous_messages
+                    }
+                })
 
             sio.on('connect', on_connect)
             sio.on('query_progress', on_query_progress)
@@ -193,7 +240,10 @@ def messages():
                         })
                 
             except Exception as e:
-                logging.error(f"Error during socket connection: {str(e)}")
+                logger.error(f"Error during socket connection", extra={
+                    'event_type': 'connection_error',
+                    'data': {'error': str(e)}
+                })
                 yield create_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
@@ -214,7 +264,10 @@ def messages():
         return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
-        logging.error(f"Request error: {str(e)}")
+        logger.error(f"Request error", extra={
+            'event_type': 'request_error',
+            'data': {'error': str(e)}
+        })
         log_request(request.remote_addr, request.path, 400)
         return jsonify({"error": str(e)}), 400
 
@@ -238,7 +291,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
                 response_event.set()
 
         def on_connect():
-            logging.info("Connected to Perplexity AI")
+            logger.info("Connected to Perplexity AI (non-stream)", extra={'event_type': 'connection_established_non_stream'})
             emit_data = {
                 "version": "2.9",
                 "source": "default",
@@ -278,7 +331,10 @@ def handle_non_stream(previous_messages, msg_id, model, input_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)}")
+        logger.error(f"Error during non-stream socket connection", extra={
+            'event_type': 'non_stream_connection_error',
+            'data': {'error': str(e)}
+        })
         return jsonify({"error": str(e)}), 500
     finally:
         if sio.connected:
@@ -291,13 +347,20 @@ def not_found(error):
 
 @app.errorhandler(500)
 def server_error(error):
-    logging.error(f"Server error: {str(error)}")
+    logger.error(f"Server error", extra={
+        'event_type': 'server_error',
+        'data': {'error': str(error)}
+    })
     log_request(request.remote_addr, request.path, 500)
     return "Something broke!", 500
 
 if __name__ == '__main__':
+    setup_logging()
     port = int(os.environ.get('PORT', 8081))
-    logging.info(f"Perplexity proxy listening on port {port}")
+    logger.info(f"Perplexity proxy starting", extra={
+        'event_type': 'server_start',
+        'data': {'port': port}
+    })
     if not API_KEY:
-        logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
-    app.run(host='0.0.0.0', port=port)
+        logger.warning("PPLX_KEY environment variable is not set", extra={'event_type': 'config_warning'})
+    app.run(host='0.0.0.0', port=port)