zhangqian il y a 1 an
Parent
commit
612d51a591
2 fichiers modifiés avec 19 ajouts et 14 suppressions
  1. 13 11
      app/api/chat.py
  2. 6 3
      app/service/ragflow.py

+ 13 - 11
app/api/chat.py

@@ -47,17 +47,19 @@ async def handle_client(websocket: WebSocket,
                     message = await websocket.receive_json()
                     print(f"Received from client {chat_id}: {message}")
                     async for rag_response in ragflow_service.chat(token, chat_id, message["chatHistory"]):
-                        print(f"Received from ragflow: {rag_response}")
-                        json_str = rag_response[5:].strip()
-                        json_data = json.loads(json_str)
-                        if json_data.get("data") is not True:
-                            answer = json_data.get("data", {}).get("answer", "")
-                            result = {"message": answer, "type": "stream"}
-                        else:
-                            result = {"message": "", "type": "close"}
-                        await websocket.send_json(result)
-                        print(f"Forwarded to client {chat_id}: {result}")
-
+                        try:
+                            print(f"Received from ragflow: {rag_response}")
+                            json_str = rag_response[5:].strip()
+                            json_data = json.loads(json_str)
+                            if json_data.get("data") is not True:
+                                answer = json_data.get("data", {}).get("answer", "")
+                                result = {"message": answer, "type": "stream"}
+                            else:
+                                result = {"message": "", "type": "close"}
+                            await websocket.send_json(result)
+                            print(f"Forwarded to client {chat_id}: {result}")
+                        except Exception as e:
+                            print(f"Error forwarding message to ragflow: {e}")
             # 启动任务处理客户端消息
             tasks = [
                 asyncio.create_task(forward_to_ragflow())

+ 6 - 3
app/service/ragflow.py

@@ -51,8 +51,11 @@ class RagflowService:
                 # 检查响应状态码
                 if response.status_code == 200:
                     # 流式读取响应
-                    async for answer in response.aiter_text():
-                        yield answer
+                    try:
+                        async for answer in response.aiter_text():
+                            yield answer
+                    except GeneratorExit as e:
+                        print(e)
+                        return
                 else:
                     yield f"Error: {response.status_code}"
-