Update app.py

app.py

@@ -86,29 +86,20 @@ def messages():
 
     try:
         json_body = request.json
-        if not json_body.get('stream', False):
-            log_request(request.remote_addr, request.path, 200)
-            return jsonify({
-                "id": str(uuid.uuid4()),
-                "content": [
-                    {"text": "Please turn on streaming."},
-                    {"id": "string", "name": "string", "input": {}}
-                ],
-                "model": "string",
-                "stop_reason": "end_turn",
-                "stop_sequence": "string",
-                "usage": {"input_tokens": 0, "output_tokens": 0}
-            })
-        
+        stream = json_body.get('stream', True)  # 默认为True
+        previous_messages = "\n\n".join([msg['content'] for msg in json_body['messages']])
+        msg_id = str(uuid.uuid4())
+        response_event = Event()
+        response_text = []
+
+        if not stream:
+            # 处理 stream 为 false 的情况
+            return handle_non_stream(previous_messages, msg_id)
+
         # 记录日志：此时请求上下文仍然有效
         log_request(request.remote_addr, request.path, 200)
 
         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 = []
-
             yield create_event("message_start", {
                 "type": "message_start",
                 "message": {
@@ -214,6 +205,86 @@ def messages():
         log_request(request.remote_addr, request.path, 400)
         return jsonify({"error": str(e)}), 400
 
+def handle_non_stream(previous_messages, msg_id):
+    """
+    处理 stream 为 false 的情况，返回完整的响应。
+    """
+    try:
+        response_event = Event()
+        response_text = []
+
+        def on_connect():
+            logging.info("Connected to Perplexity AI")
+            emit_data = {
+                "version": "2.9",
+                "source": "default",
+                "attachments": [],
+                "language": "en-GB",
+                "timezone": "Europe/London",
+                "search_focus": "writing",
+                "frontend_uuid": str(uuid.uuid4()),
+                "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():
+            logging.info("Disconnected from Perplexity AI")
+            response_event.set()
+
+        def on_connect_error(data):
+            logging.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.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
+        
+        # 等待响应完成
+        response_event.wait(timeout=30)
+
+        # 生成完整的响应
+        full_response = {
+            "id": msg_id,
+            "content": [{"text": ''.join(response_text)}],  # 合并所有文本块
+            "model": "claude-3-opus-20240229",
+            "stop_reason": "end_turn",
+            "stop_sequence": None,
+            "usage": {
+                "input_tokens": 8,
+                "output_tokens": len(''.join(response_text)),
+            },
+        }
+        return jsonify(full_response)
+
+    except Exception as e:
+        logging.error(f"Error during socket connection: {str(e)}")
+        return jsonify({"error": str(e)}), 500
+    finally:
+        if sio.connected:
+            sio.disconnect()
+
 @app.errorhandler(404)
 def not_found(error):
     log_request(request.remote_addr, request.path, 404)