From 39eba302f3887d83ce0584b7e79440a92f9037c2 Mon Sep 17 00:00:00 2001 From: zhangqian <zhangqian@123.com> Date: 星期五, 18 十月 2024 11:23:19 +0800 Subject: [PATCH] ragflow返回的消息改个类型 --- app/api/chat.py | 48 +++++++++++++++++++++++++++++++----------------- 1 files changed, 31 insertions(+), 17 deletions(-) diff --git a/app/api/chat.py b/app/api/chat.py index 458581a..07a05da 100644 --- a/app/api/chat.py +++ b/app/api/chat.py @@ -46,37 +46,52 @@ message = await websocket.receive_json() print(f"Received from client {chat_id}: {message}") chat_history = message.get('chatHistory', []) + message["role"] = "user" if len(chat_history) == 0: - - chat_history = await ragflow_service.set_session(token, agent_id, message["message"], chat_id, True) + chat_history = await ragflow_service.get_session_history(token, chat_id) if len(chat_history) == 0: - result = {"message": "鍐呴儴閿欒锛氬垱寤轰細璇濆け璐�", "type": "close"} - await websocket.send_json(result) + 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}") if rag_response[:5] == "data:": # 濡傛灉鏄紝鍒欐埅鍙栨帀鍓�5涓瓧绗︼紝骞跺幓闄ら灏剧┖鐧界 - json_str = rag_response[5:].strip() + text = rag_response[5:].strip() else: # 鍚﹀垯锛屼繚鎸佸師鏍� - json_str = rag_response - 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: # 姝e父杈撳嚭 - answer = data.get("answer", "") - result = {"message": answer, "type": "stream"} + 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: # 姝e父杈撳嚭 + answer = data.get("answer", "") + result = {"message": answer, "type": "message"} + except json.JSONDecodeError: + result = {"message": text, "type": "message"} 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 process message of ragflow: {e}") + # 鍚姩浠诲姟澶勭悊瀹㈡埛绔秷鎭� tasks = [ asyncio.create_task(forward_to_ragflow()) @@ -145,4 +160,3 @@ else: ret = {"message": "Agent not found", "type": "close"} await websocket.send_json(ret) - -- Gitblit v1.8.0