派生自 rag-gateway
app/api/chat.py
@@ -29,40 +29,69 @@ agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first() if not agent: ret = {"message": "Agent not found", "type": "close"} return websocket.send_json(ret) await websocket.send_json(ret) return agent_type = agent.agent_type if chat_id == "" or chat_id == "0": ret = {"message": "Chat ID not found", "type": "close"} return websocket.send_json(ret) await websocket.send_json(ret) return if agent_type == AgentType.RAGFLOW: ragflow_service = RagflowService(settings.ragflow_base_url) ragflow_service = RagflowService(settings.fwr_base_url) token = get_ragflow_token(db, current_user.id) try: async def forward_to_ragflow(): while True: 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"]): chat_history = message.get('chatHistory', []) message["role"] = "user" if len(chat_history) == 0: chat_history = await ragflow_service.get_session_history(token, chat_id) if len(chat_history) == 0: chat_history = await ragflow_service.set_session(token, agent_id, message, chat_id, True) if len(chat_history) == 0: result = {"message": "内部错误：创建会话失败", "type": "close"} await websocket.send_json(result) await websocket.close() return else: chat_history.append({ "content": message["message"], "doc_ids": message.get("doc_ids", []), "role": "user" }) async for rag_response in ragflow_service.chat(token, chat_id, chat_history): 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}") if rag_response[:5] == "data:": # 如果是，则截取掉前5个字符，并去除首尾空白符 text = rag_response[5:].strip() else: # 否则，保持原样 text = rag_response try: json_data = json.loads(text) 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": "message"} else: # 正常输出 answer = data.get("answer", "") result = {"message": answer, "type": "message"} await websocket.send_json(result) except json.JSONDecodeError: print(f"Error decode ragflow response: {text}") pass except Exception as e: result = {"message": f"内部错误： {e}", "type": "close"} await websocket.send_json(result) print(f"Error process message of ragflow: {e}") # 启动任务处理客户端消息 tasks = [ asyncio.create_task(forward_to_ragflow()) @@ -73,7 +102,7 @@ elif agent_type == AgentType.BISHENG: token = get_bisheng_token(db, current_user.id) service_uri = f"{settings.bisheng_websocket_url}/api/v1/assistant/chat/{agent_id}?t=&chat_id={chat_id}" service_uri = f"{settings.sgb_websocket_url}/api/v1/assistant/chat/{agent_id}?t=&chat_id={chat_id}" headers = {'cookie': f"access_token_cookie={token};"} async with websockets.connect(service_uri, extra_headers=headers) as service_websocket: @@ -130,5 +159,4 @@ print(f"Client {chat_id} disconnected") else: ret = {"message": "Agent not found", "type": "close"} return websocket.send_json(ret) await websocket.send_json(ret)
