Update app.py

app.py

@@ -66,7 +66,7 @@ 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 jsonify({"error": "Invalid API key"}), 401
+        return create_json_response({"error": "Invalid API key"}), 401
     return None
 
 def normalize_content(content):
@@ -95,10 +95,16 @@ def calculate_tokens(text):
         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')
+
 @app.route('/')
 def root():
     log_request(request.remote_addr, request.path, 200)
-    return jsonify({
+    return create_json_response({
         "message": "Welcome to the Perplexity AI Proxy API",
         "endpoints": {
             "/ai/v1/messages": {
@@ -141,73 +147,72 @@ def messages():
         log_request(request.remote_addr, request.path, 200)
 
         def generate():
-            yield create_event("message_start", {
-                "type": "message_start",
-                "message": {
-                    "id": msg_id,
-                    "type": "message",
-                    "role": "assistant",
-                    "content": [],
-                    "model": model,
-                    "stop_reason": None,
-                    "stop_sequence": None,
-                    "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():
-                app_logger.info("Connected to Perplexity AI")
-                emit_data = {
-                    "version": "2.9",
-                    "source": "default",
-                    "attachments": [],
-                    "language": "en-GB",
-                    "timezone": "Europe/London",
-                    "mode": "concise",
-                    "is_related_query": False,
-                    "is_default_related_query": False,
-                    "visitor_id": str(uuid.uuid4()),
-                    "frontend_context_uuid": str(uuid.uuid4()),
-                    "prompt_source": "user",
-                    "query_source": "home"
-                }
-                sio.emit('perplexity_ask', (previous_messages, emit_data))
-
-            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)
-
-                if data.get('final', False):
+            try:
+                yield create_event("message_start", {
+                    "type": "message_start",
+                    "message": {
+                        "id": msg_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [],
+                        "model": model,
+                        "stop_reason": None,
+                        "stop_sequence": None,
+                        "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():
+                    app_logger.info("Connected to Perplexity AI")
+                    emit_data = {
+                        "version": "2.9",
+                        "source": "default",
+                        "attachments": [],
+                        "language": "en-GB",
+                        "timezone": "Europe/London",
+                        "mode": "concise",
+                        "is_related_query": False,
+                        "is_default_related_query": False,
+                        "visitor_id": str(uuid.uuid4()),
+                        "frontend_context_uuid": str(uuid.uuid4()),
+                        "prompt_source": "user",
+                        "query_source": "home"
+                    }
+                    sio.emit('perplexity_ask', (previous_messages, emit_data))
+
+                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)
+
+                    if data.get('final', False):
+                        response_event.set()
+
+                def on_disconnect():
+                    app_logger.info("Disconnected from Perplexity AI")
                     response_event.set()
 
-            def on_disconnect():
-                app_logger.info("Disconnected from Perplexity AI")
-                response_event.set()
-
-            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()
+                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.on('connect', on_connect)
+                sio.on('query_progress', on_query_progress)
+                sio.on('disconnect', on_disconnect)
+                sio.on('connect_error', on_connect_error)
 
-            try:
                 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",
@@ -215,33 +220,29 @@ def messages():
                             "delta": {"type": "text_delta", "text": chunk},
                         })
                 
-            except Exception as e:
-                app_logger.error(f"Error during socket connection: {str(e)}")
-                yield create_event("content_block_delta", {
-                    "type": "content_block_delta",
-                    "index": 0,
-                    "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
+                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},
                 })
+                yield create_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)})
             finally:
                 if sio.connected:
                     sio.disconnect()
 
-            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},
-            })
-            yield create_event("message_stop", {"type": "message_stop"})
-
         return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
         app_logger.error(f"Request error: {str(e)}")
         log_request(request.remote_addr, request.path, 400)
-        return jsonify({"error": str(e)}), 400
+        return create_json_response({"error": str(e)}), 400
 
 def handle_non_stream(previous_messages, msg_id, model, input_tokens):
     try:
@@ -310,13 +311,11 @@ def handle_non_stream(previous_messages, msg_id, model, input_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')
+        return create_json_response(full_response)
 
     except Exception as e:
         app_logger.error(f"Error during socket connection: {str(e)}")
-        return jsonify({"error": str(e)}), 500
+        return create_json_response({"error": str(e)}), 500
     finally:
         if sio.connected:
             sio.disconnect()
@@ -324,18 +323,20 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
 @app.errorhandler(404)
 def not_found(error):
     log_request(request.remote_addr, request.path, 404)
-    return "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 "Something broke!", 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)
-    return f"event: {event}\ndata: {data}\n\n"
+    event_str = f"event: {event}\ndata: {data}\n\n"
+    app_logger.info(f"Sending SSE event: {event_str}")
+    return event_str
 
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 8081))