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, copy_current_request_context
+from flask import Flask, request, Response, jsonify
 import socketio
 import requests
 import logging
@@ -100,13 +100,11 @@ def messages():
                 "usage": {"input_tokens": 0, "output_tokens": 0}
             })
         
-        @stream_with_context
         def generate():
             previous_messages = "\n\n".join([msg['content'] for msg in json_body['messages']])
             msg_id = str(uuid.uuid4())
             response_event = Event()
             response_text = []
-            client_disconnected = False
 
             yield create_event("message_start", {
                 "type": "message_start",
@@ -150,9 +148,16 @@ def messages():
                     text = json.loads(data['text'])
                     chunks = text.get('chunks', [])
                     for chunk in chunks:
-                        response_text.append(chunk)
-                if data.get('final', False):
-                    response_event.set()
+                        decoded_chunk = chunk.encode('utf-8').decode('unicode_escape')
+                        response_text.append(decoded_chunk)
+                        yield create_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": decoded_chunk},
+                        })
+
+            def on_query_complete(data):
+                response_event.set()
 
             def on_disconnect():
                 logging.info("Disconnected from Perplexity AI")
@@ -160,31 +165,22 @@ def messages():
 
             def on_connect_error(data):
                 logging.error(f"Connection error: {data}")
-                response_text.append(f"Error connecting to Perplexity AI: {data}")
+                yield create_event("content_block_delta", {
+                    "type": "content_block_delta",
+                    "index": 0,
+                    "delta": {"type": "text_delta", "text": f"Error connecting to Perplexity AI: {data}"},
+                })
                 response_event.set()
 
             sio.on('connect', on_connect)
             sio.on('query_progress', on_query_progress)
+            sio.on('query_complete', on_query_complete)
             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() and not client_disconnected:
-                    sio.sleep(0.1)
-                    while response_text:
-                        chunk = response_text.pop(0)
-                        try:
-                            yield create_event("content_block_delta", {
-                                "type": "content_block_delta",
-                                "index": 0,
-                                "delta": {"type": "text_delta", "text": chunk},
-                            })
-                        except GeneratorExit:
-                            client_disconnected = True
-                            break
-                
+                response_event.wait(timeout=30)  # 等待响应，最多30秒
             except Exception as e:
                 logging.error(f"Error during socket connection: {str(e)}")
                 yield create_event("content_block_delta", {
@@ -196,21 +192,16 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            if not client_disconnected:
-                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": {"output_tokens": len(''.join(response_text))},
-                })
-                yield create_event("message_stop", {"type": "message_stop"})
-            
+            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": {"output_tokens": len(''.join(response_text))},
+            })
+            yield create_event("message_stop", {"type": "message_stop"})
             log_request(request.remote_addr, request.path, 200)
 
-        return Response(generate(), content_type='text/event-stream', headers={
-            'Cache-Control': 'no-cache',
-            'X-Accel-Buffering': 'no'
-        })
+        return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
         logging.error(f"Request error: {str(e)}")
@@ -230,8 +221,8 @@ def server_error(error):
 
 def create_event(event, data):
     if isinstance(data, dict):
-        data = json.dumps(data, ensure_ascii=False)
-    return f"event: {event}"
+        data = json.dumps(data)
+    return f"event: {event}\ndata: {data}\n\n"
 
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 8081))