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