zhangqian
2024-10-14 5d535b127cffb34affcb5d2a15fb3ddd4aa7e5ba
app/api/chat.py
@@ -47,17 +47,25 @@
                    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)
                            data = json_data.get("data")
                            if data is True:  # 完成输出
                                result = {"message": "", "type": "close"}
                            elif data is None:  # 发生错误
                                answer = json_data.get("retmsg", json_data.get("retcode"))
                                result = {"message": "内部错误:" + answer, "type": "stream"}
                            else:  # 正常输出
                                answer = json_data.get("data", {}).get("answer", "")
                                result = {"message": answer, "type": "stream"}
                            await websocket.send_json(result)
                            print(f"Forwarded to client {chat_id}: {result}")
                        except Exception as e:
                            result = {"message": f"内部错误: {e}", "type": "close"}
                            await websocket.send_json(result)
                            print(f"Error forwarding message to ragflow: {e}")
            # 启动任务处理客户端消息
            tasks = [
                asyncio.create_task(forward_to_ragflow())