Update app.py

app.py

@@ -1,37 +1,15 @@
 import os
 import json
 import uuid
+import re
 from datetime import datetime
 from flask import Flask, request, Response, jsonify
 import socketio
 import requests
 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__)
-logger = logging.getLogger(__name__)
 
 # 从环境变量中获取API密钥
 API_KEY = os.environ.get('PPLX_KEY')
@@ -54,7 +32,7 @@ sio = socketio.Client(http_session=transport, logger=False, engineio_logger=Fals
 
 # 连接选项
 connect_opts = {
-    'transports': ['websocket', 'polling'],  # 允许回退到轮询
+    'transports': ['websocket', 'polling'],
 }
 
 # 其他选项
@@ -68,10 +46,38 @@ sio_opts = {
     }
 }
 
+class CustomFormatter(logging.Formatter):
+    def format(self, record):
+        log_data = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "message": self.remove_ansi_escape(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 remove_ansi_escape(self, text):
+        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+        return ansi_escape.sub('', text)
+
+def setup_logging():
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler()
+    handler.setFormatter(CustomFormatter())
+    logger.addHandler(handler)
+
+logger = logging.getLogger(__name__)
+
 def log_request(ip, route, status):
-    logger.info(f"HTTP Request", extra={
-        'event_type': 'http_request',
+    timestamp = datetime.now().isoformat()
+    logger.info(f"Request received", extra={
+        'event_type': 'request',
         'data': {
+            'timestamp': timestamp,
             'ip': ip,
             'route': route,
             'status': status
@@ -156,7 +162,7 @@ def messages():
         def generate():
             nonlocal total_output_tokens
 
-            yield create_event("message_start", {
+            start_event = create_event("message_start", {
                 "type": "message_start",
                 "message": {
                     "id": msg_id,
@@ -169,8 +175,25 @@ def messages():
                     "usage": {"input_tokens": input_tokens, "output_tokens": total_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"})
+            logger.info("Sending message_start event", extra={
+                'event_type': 'message_start',
+                'data': {'content': start_event}
+            })
+            yield start_event
+
+            block_start_event = create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+            logger.info("Sending content_block_start event", extra={
+                'event_type': 'content_block_start',
+                'data': {'content': block_start_event}
+            })
+            yield block_start_event
+
+            ping_event = create_event("ping", {"type": "ping"})
+            logger.info("Sending ping event", extra={
+                'event_type': 'ping',
+                'data': {'content': ping_event}
+            })
+            yield ping_event
 
             def on_query_progress(data):
                 nonlocal total_output_tokens, response_text
@@ -181,22 +204,23 @@ 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',
+                        delta_event = create_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": chunk},
+                        })
+                        logger.info("Sending content_block_delta event", extra={
+                            'event_type': 'content_block_delta',
                             'data': {
                                 'chunk': chunk,
                                 'tokens': chunk_tokens,
-                                'total_tokens': total_output_tokens
+                                'total_tokens': total_output_tokens,
+                                'content': delta_event
                             }
                         })
+                        yield delta_event
 
                 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():
@@ -231,16 +255,9 @@ def messages():
                 
                 while not response_event.is_set():
                     sio.sleep(0.1)
-                    while response_text:
-                        chunk = response_text.pop(0)
-                        yield create_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
                 
             except Exception as e:
-                logger.error(f"Error during socket connection", extra={
+                logger.error(f"Error during socket connection: {str(e)}", extra={
                     'event_type': 'connection_error',
                     'data': {'error': str(e)}
                 })
@@ -253,18 +270,35 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            yield create_event("message_delta", {
+            stop_event = create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+            logger.info("Sending content_block_stop event", extra={
+                'event_type': 'content_block_stop',
+                'data': {'content': stop_event}
+            })
+            yield stop_event
+
+            message_delta_event = create_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                 "usage": {"output_tokens": total_output_tokens},
             })
-            yield create_event("message_stop", {"type": "message_stop"})
+            logger.info("Sending message_delta event", extra={
+                'event_type': 'message_delta',
+                'data': {'content': message_delta_event}
+            })
+            yield message_delta_event
+
+            message_stop_event = create_event("message_stop", {"type": "message_stop"})
+            logger.info("Sending message_stop event", extra={
+                'event_type': 'message_stop',
+                'data': {'content': message_stop_event}
+            })
+            yield message_stop_event
 
         return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
-        logger.error(f"Request error", extra={
+        logger.error(f"Request error: {str(e)}", extra={
             'event_type': 'request_error',
             'data': {'error': str(e)}
         })
@@ -328,11 +362,15 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
                 "output_tokens": total_output_tokens,
             },
         }
+        logger.info("Sending non-stream response", extra={
+            'event_type': 'non_stream_response',
+            'data': {'content': full_response}
+        })
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
 
     except Exception as e:
-        logger.error(f"Error during non-stream socket connection", extra={
-            'event_type': 'non_stream_connection_error',
+        logger.error(f"Error during socket connection: {str(e)}", extra={
+            'event_type': 'connection_error_non_stream',
             'data': {'error': str(e)}
         })
         return jsonify({"error": str(e)}), 500
@@ -347,7 +385,7 @@ def not_found(error):
 
 @app.errorhandler(500)
 def server_error(error):
-    logger.error(f"Server error", extra={
+    logger.error(f"Server error: {str(error)}", extra={
         'event_type': 'server_error',
         'data': {'error': str(error)}
     })
@@ -357,7 +395,7 @@ def server_error(error):
 if __name__ == '__main__':
     setup_logging()
     port = int(os.environ.get('PORT', 8081))
-    logger.info(f"Perplexity proxy starting", extra={
+    logger.info("Perplexity proxy starting", extra={
         'event_type': 'server_start',
         'data': {'port': port}
     })