Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import json
 import uuid
 from datetime import datetime
-from flask import Flask, request, Response, jsonify
+from flask import Flask, request, Response, jsonify, stream_with_context
 import socketio
 import requests
 import logging
@@ -31,7 +31,7 @@ class CustomFormatter(logging.Formatter):
 
 def setup_logging():
     logger = logging.getLogger()
-    logger.setLevel(logging.INFO)
+    logger.setLevel(logging.DEBUG)  # 设置为DEBUG级别以捕获更多信息
     handler = logging.StreamHandler()
     handler.setFormatter(CustomFormatter())
     logger.addHandler(handler)
@@ -164,16 +164,15 @@ def messages():
         def generate():
             nonlocal total_output_tokens
 
+            yield 'retry: 1000\n\n'  # 设置重试间隔
+
             def send_event(event_type, data):
                 event = create_event(event_type, data)
-                logger.info(f"Sending {event_type} event", extra={
-                    'event_type': event_type,
-                    'data': {'content': event}
-                })
-                yield event
+                logger.debug(f"Sending event: {event_type}")
+                return event
 
-            # Send initial events
-            yield from send_event("message_start", {
+            # 发送初始事件
+            yield send_event("message_start", {
                 "type": "message_start",
                 "message": {
                     "id": msg_id,
@@ -186,8 +185,8 @@ def messages():
                     "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
                 },
             })
-            yield from send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
-            yield from send_event("ping", {"type": "ping"})
+            yield send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+            yield send_event("ping", {"type": "ping"})
 
             def on_query_progress(data):
                 nonlocal total_output_tokens, response_text
@@ -258,14 +257,14 @@ def messages():
                     sio.sleep(0.1)
                     while response_text:
                         chunk = response_text.pop(0)
-                        yield from send_event("content_block_delta", {
+                        yield send_event("content_block_delta", {
                             "type": "content_block_delta",
                             "index": 0,
                             "delta": {"type": "text_delta", "text": chunk},
                         })
                 
                 if timeout_event.is_set():
-                    yield from send_event("content_block_delta", {
+                    yield send_event("content_block_delta", {
                         "type": "content_block_delta",
                         "index": 0,
                         "delta": {"type": "text_delta", "text": "Request timed out"},
@@ -273,7 +272,7 @@ def messages():
                 
             except Exception as e:
                 logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
-                yield from send_event("content_block_delta", {
+                yield send_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
                     "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
@@ -283,16 +282,21 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            # Send final events
-            yield from send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            yield from send_event("message_delta", {
+            # 发送最终事件
+            yield send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+            yield send_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                 "usage": {"output_tokens": total_output_tokens},
             })
-            yield from send_event("message_stop", {"type": "message_stop"})
+            yield send_event("message_stop", {"type": "message_stop"})
 
-        return Response(generate(), content_type='text/event-stream')
+        headers = {
+            'Content-Type': 'text/event-stream',
+            'Cache-Control': 'no-cache',
+            'X-Accel-Buffering': 'no',
+        }
+        return Response(stream_with_context(generate()), headers=headers)
 
     except Exception as e:
         logger.error(f"Request error: {str(e)}", exc_info=True)
@@ -394,4 +398,5 @@ if __name__ == '__main__':
     })
     if not API_KEY:
         logger.warning("PPLX_KEY environment variable is not set", extra={'event_type': 'config_warning'})
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port, debug=False)
+