From ac9373fd6f0db53c983d87c4790cb627e82759e9 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期二, 25 二月 2025 17:25:39 +0800 Subject: [PATCH] merge --- app/api/chat.py | 71 +++++++++++++++++++++++++++++++---- 1 files changed, 63 insertions(+), 8 deletions(-) diff --git a/app/api/chat.py b/app/api/chat.py index 6e54cb9..241815a 100644 --- a/app/api/chat.py +++ b/app/api/chat.py @@ -11,7 +11,8 @@ from app.api import get_current_user_websocket from app.config.config import settings from app.config.const import IMAGE_TO_TEXT, DOCUMENT_TO_REPORT, DOCUMENT_TO_CLEANING, DOCUMENT_IA_QUESTIONS, \ - DOCUMENT_TO_REPORT_TITLE, DOCUMENT_TO_TITLE, DOCUMENT_TO_PAPER, DOCUMENT_IA_QUESTIONS_DS + DOCUMENT_TO_REPORT_TITLE, DOCUMENT_TO_TITLE, DOCUMENT_TO_PAPER, DOCUMENT_IA_QUESTIONS_DS, \ + DOCUMENT_IA_QUESTIONS_EQUIPMENT from app.models import MenuCapacityModel from app.models.agent_model import AgentModel, AgentType from app.models.base_model import get_db @@ -39,6 +40,7 @@ print(f"Client {agent_id} connected") agent = db.query(MenuCapacityModel).filter(MenuCapacityModel.chat_id == agent_id).first() if not agent: + print("Agent not found") agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first() agent_type = agent.agent_type chat_type = agent.type @@ -808,13 +810,17 @@ result = {"message": f"鍐呴儴閿欒锛� {e2}", "type": "close"} await websocket.send_json(result) print(f"Error process message of ragflow: {e2}") - elif chat_type == "documentIa" or chat_type == "documentIaDs": + elif chat_type == "documentIa" or chat_type == "documentIaDs" or chat_type == "documentIaEq": token_dict = { "documentIa": DOCUMENT_IA_QUESTIONS, "documentIaDs": DOCUMENT_IA_QUESTIONS_DS, + "documentIaEq": DOCUMENT_IA_QUESTIONS_EQUIPMENT, } token = DfTokenDao(db).get_token_by_id(token_dict[chat_type]) # print(token) + node_list = [] + workflow_list = [] + workflow_dict = {} if not token: await websocket.send_json({"message": "Invalid token", "type": "error"}) @@ -867,8 +873,15 @@ complete_response += rag_response try: data = json.loads(complete_response) - if data.get("event") == "node_started" or data.get( - "event") == "node_finished": # "event": "message_end" + node_data = deepcopy(data) + if "data" in node_data: + if "outputs" in node_data["data"]: + node_data["data"]["outputs"] = {} + if "inputs" in node_data["data"]: + node_data["data"]["inputs"] = {} + # print(node_data) + node_list.append(node_data) + if data.get("event") == "node_started": if "data" not in data or not data["data"]: # 淇℃伅杩囨护 logger.error("闈炴硶鏁版嵁--------------------") logger.error(data) @@ -885,12 +898,43 @@ else: message = answer.get("title", "") - result = {"message": message, "type": "system"} - # continue + result = {"message": message, "type": "system", "workflow": {"node_data": workflow_list}} + elif data.get("event") == "node_finished": + workflow_list.append({ + "title": data.get("data", {}).get("title", ""), + "status": data.get("data", {}).get("status", ""), + "created_at": data.get("data", {}).get("created_at", 0), + "finished_at": data.get("data", {}).get("finished_at", 0), + "node_type": data.get("data", {}).get("node_type", 0), + "elapsed_time": data.get("data", {}).get("elapsed_time", 0), + "error": data.get("data", {}).get("error", ""), + }) + if "data" not in data or not data["data"]: # 淇℃伅杩囨护 + logger.error("闈炴硶鏁版嵁--------------------") + logger.error(data) + continue + else: # 姝e父杈撳嚭 + answer = data.get("data", "") + if isinstance(answer, str): + logger.error("----------------鏈煡鏁版嵁--------------------") + logger.error(data) + continue + elif isinstance(answer, dict): + if answer.get("status") == "failed": + message = answer.get("error", "") + else: + message = answer.get("title", "") + + result = {"message": message, "type": "system", + "workflow": {"node_data": workflow_list}} elif data.get("event") == "message": # "event": "message_end" # 姝e父杈撳嚭 answer = data.get("answer", "") - result = {"message": answer, "type": "stream"} + result = {"message": answer, "type": "stream", "workflow": {"node_data": workflow_list}} + elif data.get("event") == "error": # "event": "message_end" + # 姝e父杈撳嚭 + answer = data.get("message", "") + result = {"message": answer, "type": "system", "workflow": {"node_data": workflow_list}} elif data.get("event") == "workflow_finished": answer = data.get("data", "") if isinstance(answer, str): @@ -913,16 +957,27 @@ message={"role": "assistant", "content": { "answer": message, + "node_list": node_list, "download_url": download_url}}, conversation_id=data.get( "conversation_id")) + node_list = [] except Exception as e: logger.error("淇濆瓨dify鐨勪細璇濆紓甯革紒") logger.error(e) # await websocket.send_json(result) + workflow_dict = { + "node_data": workflow_list, + "total_tokens": data.get("data", {}).get("total_tokens", 0), + "created_at": data.get("data", {}).get("created_at", 0), + "finished_at": data.get("data", {}).get("finished_at", 0), + "status": data.get("data", {}).get("status", ""), + "error": data.get("data", {}).get("error", ""), + "elapsed_time": data.get("data", {}).get("elapsed_time", 0) + } continue elif data.get("event") == "message_end": - result = {"message": "", "type": "close"} + result = {"message": "", "type": "close","workflow": workflow_dict} else: continue -- Gitblit v1.8.0