From 72a8a0a1ad6b79b8e9fb2facef121f9b5d584666 Mon Sep 17 00:00:00 2001
From: xuyonghao <898441624@qq.com>
Date: 星期六, 08 二月 2025 10:56:30 +0800
Subject: [PATCH] 报表合并历史记录功能

---
 app/api/chat.py |  308 ++++++++++++++++++++++++++++++++++----------------
 1 files changed, 207 insertions(+), 101 deletions(-)

diff --git a/app/api/chat.py b/app/api/chat.py
index 2993b2a..5914353 100644
--- a/app/api/chat.py
+++ b/app/api/chat.py
@@ -1,6 +1,7 @@
 import json
 import re
 import uuid
+from copy import deepcopy
 
 from fastapi import WebSocket, WebSocketDisconnect, APIRouter, Depends
 import asyncio
@@ -423,9 +424,7 @@
                                 print(f"Error process message of ragflow: {e2}")
                 elif agent.type == "reportWorkflow":
 
-                    token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_CLEANING)
-                    if not token:
-                        await websocket.send_json({"message": "Invalid token document_to_cleaning", "type": "error"})
+
                     while True:
                         receive_message = await websocket.receive_json()
                         print(f"Received from client {chat_id}: {receive_message}")
@@ -451,7 +450,9 @@
                                 title if title else title_query,
                                 agent_id,
                                 AgentType.DIFY,
-                                current_user.id
+                                current_user.id,
+                                {"role": "user", "content": title if title else title_query, "type": workflow_type, "is_clean":is_clean},
+                                workflow_type
                             )
                             conversation_id = session.conversation_id
                         except Exception as e:
@@ -475,12 +476,16 @@
                                     "upload_file_id": ""
                                 })
                         inputs_list = []
-                        token_list = []
+                        is_next = 0
                         if workflow_type == 1:
+                            token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_CLEANING)
+                            if not token:
+                                await websocket.send_json(
+                                    {"message": "Invalid token document_to_cleaning", "type": "error"})
                             inputs["input_files"] = files
-                            inputs_list.append(inputs)
-                            token_list.append(token)
-                        elif workflow_type == 2 and is_clean == 0:
+                            inputs["Completion_of_main_indicators"] = title
+                            inputs_list.append({"inputs": inputs, "token": token, "workflow_type": workflow_type})
+                        elif workflow_type == 2:
                             inputs["file_list"] = files
                             inputs["Completion_of_main_indicators"] = title
                             inputs["sub_titles"] = sub_titles
@@ -488,9 +493,8 @@
                             if not token:
                                 await websocket.send_json(
                                     {"message": "Invalid token document_to_report", "type": "error"})
-                            inputs_list.append(inputs)
-                            token_list.append(token)
-                        elif workflow_type == 3:
+                            inputs_list.append({"inputs": inputs, "token": token, "workflow_type": workflow_type})
+                        elif workflow_type == 3 and is_clean == 0 and tokens < max_token:
                             inputs["file_list"] = files
                             inputs["number_of_title"] = title_number
                             inputs["title_style"] = title_style
@@ -498,28 +502,42 @@
                             if not token:
                                 await websocket.send_json(
                                     {"message": "Invalid token document_to_title", "type": "error"})
-                            # inputs_list.append(inputs)
-                            # token_list.append(token)
-                        elif workflow_type == 2 and is_clean == 1:
-                            # inputs["input_files"] = files
-                            inputs_list.append(inputs)
-                            token_list.append(token)
+                            inputs_list.append({"inputs": inputs, "token": token, "workflow_type": workflow_type})
+                        elif workflow_type == 3 and is_clean == 1 or tokens >= max_token:
+                            token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_CLEANING)
+                            if not token:
+                                await websocket.send_json(
+                                    {"message": "Invalid token document_to_cleaning", "type": "error"})
+                            inputs["input_files"] = files
+                            inputs["Completion_of_main_indicators"] = title
+                            inputs_list.append({"inputs": inputs, "token": token, "workflow_type": 1})
                             inputs1 = {}
-                            # inputs1["file_list"] = files
-                            inputs1["Completion_of_main_indicators"] = title
-                            inputs1["sub_titles"] = sub_titles
-                            token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_REPORT_TITLE)
+                            inputs1["file_list"] = files
+                            inputs1["number_of_title"] = title_number
+                            inputs1["title_style"] = title_style
+                            token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_TITLE)
                             if not token:
                                 await websocket.send_json(
                                     {"message": "Invalid token document_to_report", "type": "error"})
-                            inputs_list.append(inputs1)
-                            token_list.append(token)
+                            inputs_list.append({"inputs": inputs1, "token": token, "workflow_type": 3})
 
-                        complete_response = ""
-                        if workflow_type == 1 or workflow_type == 2:
-                            for inputs in inputs_list:
-                                inputs["input_files"] = files
-                                async for rag_response in dify_service.workflow(token, current_user.id, inputs):
+                        # print(inputs_list)
+                        for idx, input in enumerate(inputs_list):
+                            # print(input)
+                            if idx < len(inputs_list)-1:
+                                is_next = 1
+                            else:
+                                is_next = 0
+                            i = input["inputs"]
+                            if "file_list" in i:
+                                i["file_list"] = files
+                            # print(i)
+                            node_list = []
+                            complete_response = ""
+                            workflow_list = []
+                            workflow_dict = {}
+                            if input["workflow_type"] == 1 or input["workflow_type"] == 2:
+                                async for rag_response in dify_service.workflow(input["token"], current_user.id, i):
                                     # print(rag_response)
                                     try:
                                         if rag_response[:5] == "data:":
@@ -532,8 +550,19 @@
                                             complete_response += rag_response
                                         try:
                                             data = json.loads(complete_response)
+                                            # print(data)
+                                            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)
+
                                             complete_response = ""
-                                            if data.get("event") == "node_started" or data.get("event") == "node_finished":  # "event": "message_end"
+                                            if data.get("event") == "node_started":  # "event": "message_end"
+
                                                 if "data" not in data or not data["data"]:  # 淇℃伅杩囨护
                                                     logger.error("闈炴硶鏁版嵁--------------------")
                                                     logger.error(data)
@@ -548,13 +577,36 @@
 
                                                         message = answer.get("title", "")
 
-                                                    result = {"message": message, "type": "system"}
+
+                                                    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", ""),
+                                                })
+                                                answer = data.get("data", "")
+                                                if isinstance(answer, str):
+                                                    logger.error("----------------鏈煡鏁版嵁--------------------")
+                                                    logger.error(data)
+                                                    continue
+                                                elif isinstance(answer, dict):
+
+                                                    message = answer.get("title", "")
+                                                    if answer.get("status") == "failed":
+                                                        message = answer.get("error", "")
+                                                        result = {"message": message, "type": "system", "workflow":{"node_data": workflow_list}}
+
                                             elif data.get("event") == "workflow_finished":
                                                 answer = data.get("data", "")
                                                 if isinstance(answer, str):
                                                     logger.error("----------------鏈煡鏁版嵁--------------------")
                                                     logger.error(data)
-                                                    result = {"message": "", "type": "close", "download_url": ""}
+                                                    result = {"message": "", "type": "close", "download_url": "", "is_next": is_next}
                                                 elif isinstance(answer, dict):
                                                     download_url = ""
                                                     outputs = answer.get("outputs", {})
@@ -570,15 +622,26 @@
                                                             "url": download_url,
                                                             "upload_file_id": ""
                                                         }]
-                                                    result = {"message": message, "type": "message", "download_url": download_url}
+                                                    workflow_dict = {
+                                                        "node_data": workflow_list,
+                                                        "total_tokens": answer.get("total_tokens", 0),
+                                                        "created_at": answer.get("created_at", 0),
+                                                        "finished_at": answer.get("finished_at", 0),
+                                                        "status": answer.get("status", ""),
+                                                        "error": answer.get("error", ""),
+                                                        "elapsed_time": answer.get("elapsed_time", 0)
+                                                    }
+                                                    result = {"message": message, "type": "message", "download_url": download_url, "workflow":workflow_dict}
                                                     try:
                                                         SessionService(db).update_session(chat_id,
                                                                                           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)
@@ -587,7 +650,8 @@
                                                     except Exception as e:
                                                         logger.error(e)
                                                         logger.error("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
-                                                    result = {"message": "", "type": "close", "download_url": ""}
+
+                                                    result = {"message": "", "type": "close", "workflow": workflow_dict, "is_next": is_next, "download_url": download_url}
 
 
                                             else:
@@ -598,6 +662,7 @@
                                                 logger.error(e)
                                                 logger.error("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
                                             complete_response = ""
+
                                         except json.JSONDecodeError as e:
                                             print(f"Error decoding JSON: {e}")
                                             # print(f"Response text: {text}")
@@ -605,32 +670,61 @@
                                         result = {"message": f"鍐呴儴閿欒锛� {e2}", "type": "close"}
                                         await websocket.send_json(result)
                                         print(f"Error process message of ragflow: {e2}")
-                        elif workflow_type == 3:
-                            image_list = []
-                            # print(inputs)
-                            complete_response = ""
-                            async for rag_response in dify_service.chat(token, current_user.id, title_query, [],
-                                                                        conversation_id, inputs):
-                                print(rag_response)
-                                try:
-                                    if rag_response[:5] == "data:":
-                                        # 濡傛灉鏄紝鍒欐埅鍙栨帀鍓�5涓瓧绗︼紝骞跺幓闄ら灏剧┖鐧界
-                                        complete_response = rag_response[5:].strip()
-                                    elif "event: ping" in rag_response:
-                                        continue
-                                    else:
-                                        # 鍚﹀垯锛屼繚鎸佸師鏍�
-                                        complete_response += rag_response
+                            elif input["workflow_type"] == 3:
+                                image_list = []
+                                # print(inputs)
+                                complete_response = ""
+                                answer_str = ""
+                                async for rag_response in dify_service.chat(input["token"], current_user.id, title_query, [],
+                                                                            conversation_id, i):
+                                    # print(rag_response)
                                     try:
-                                        data = json.loads(complete_response)
-                                        complete_response = ""
-                                        if data.get("event") == "node_started" or data.get(
-                                                "event") == "node_finished":  # "event": "message_end"
-                                            if "data" not in data or not data["data"]:  # 淇℃伅杩囨护
-                                                logger.error("闈炴硶鏁版嵁--------------------")
-                                                logger.error(data)
-                                                continue
-                                            else:  # 姝e父杈撳嚭
+                                        if rag_response[:5] == "data:":
+                                            # 濡傛灉鏄紝鍒欐埅鍙栨帀鍓�5涓瓧绗︼紝骞跺幓闄ら灏剧┖鐧界
+                                            complete_response = rag_response[5:].strip()
+                                        elif "event: ping" in rag_response:
+                                            continue
+                                        else:
+                                            # 鍚﹀垯锛屼繚鎸佸師鏍�
+                                            complete_response += rag_response
+                                        try:
+                                            data = json.loads(complete_response)
+                                            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)
+                                            complete_response = ""
+                                            if data.get("event") == "node_started":  # "event": "message_end"
+                                                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):
+
+                                                        message = answer.get("title", "")
+
+                                                    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", ""),
+                                                })
+
                                                 answer = data.get("data", "")
                                                 if isinstance(answer, str):
                                                     logger.error("----------------鏈煡鏁版嵁--------------------")
@@ -639,54 +733,66 @@
                                                 elif isinstance(answer, dict):
 
                                                     message = answer.get("title", "")
+                                                    if answer.get("status") == "failed":
+                                                        message = answer.get("error", "")
+                                                    result = {"message": message, "type": "system", "workflow":{"node_data": workflow_list}}
+                                            elif data.get("event") == "message":
+                                                answer_str = data.get("answer", "")
+                                                # try:
+                                                #     msg_dict = json.loads(answer)
+                                                #     message = msg_dict.get("output",  "")
+                                                # except Exception as e:
+                                                #     print(e)
+                                                #     continue
+                                                result = {"message": answer_str, "type": "message",
+                                                          "download_url": "", "workflow": {"node_data": workflow_list}}
 
-                                                result = {"message": message, "type": "system"}
-                                        elif data.get("event") == "message":
-                                            message = data.get("answer", "")
-                                            # try:
-                                            #     msg_dict = json.loads(answer)
-                                            #     message = msg_dict.get("output",  "")
-                                            # except Exception as e:
-                                            #     print(e)
-                                            #     continue
-                                            result = {"message": message, "type": "message",
-                                                      "download_url": ""}
+                                                # try:
+                                                #     await websocket.send_json(result)
+                                                # except Exception as e:
+                                                #     logger.error(e)
+                                                #     logger.error("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
+                                            elif data.get("event") == "workflow_finished":
+                                                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)
+                                                }
+                                                try:
+                                                    SessionService(db).update_session(chat_id,
+                                                                                      message={"role": "assistant",
+                                                                                               "content": {
+                                                                                                   "answer": answer_str,
+                                                                                                   "node_list": node_list,
+                                                                                                   "download_url": ""}},
+                                                                                      conversation_id=data.get(
+                                                                                          "conversation_id"))
+                                                    node_list = []
+                                                except Exception as e:
+                                                    logger.error("淇濆瓨dify鐨勪細璇濆紓甯革紒")
+                                                    logger.error(e)
+                                            elif data.get("event") == "message_end":
+                                                result = {"message": "", "type": "close", "workflow": workflow_dict, "is_next": is_next}
+                                            else:
+                                                continue
                                             try:
-                                                SessionService(db).update_session(chat_id,
-                                                                                  message={"role": "assistant",
-                                                                                           "content": {
-                                                                                               "answer": message,
-                                                                                               "download_url": ""}},
-                                                                                  conversation_id=data.get(
-                                                                                      "conversation_id"))
+                                                await websocket.send_json(result)
                                             except Exception as e:
-                                                logger.error("淇濆瓨dify鐨勪細璇濆紓甯革紒")
                                                 logger.error(e)
-                                            # try:
-                                            #     await websocket.send_json(result)
-                                            # except Exception as e:
-                                            #     logger.error(e)
-                                            #     logger.error("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
-
-                                        elif data.get("event") == "message_end":
-                                            result = {"message": "", "type": "close", "download_url": ""}
-                                        else:
-                                            continue
-                                        try:
-                                            await websocket.send_json(result)
-                                        except Exception as e:
-                                            logger.error(e)
-                                            logger.error("dify杩斿洖瀹㈡埛绔秷鎭紓甯�!")
-                                        complete_response = ""
-                                    except json.JSONDecodeError as e:
-                                        print(f"Error decoding JSON: {e}")
-                                        # print(f"Response text: {text}")
-                                except Exception as e2:
-                                    result = {"message": f"鍐呴儴閿欒锛� {e2}", "type": "close"}
-                                    await websocket.send_json(result)
-                                    print(f"Error process message of ragflow: {e2}")
+                                                logger.error("dify杩斿洖瀹㈡埛绔秷鎭紓甯�!")
+                                            complete_response = ""
+                                        except json.JSONDecodeError as e:
+                                            print(f"Error decoding JSON: {e}")
+                                            # print(f"Response text: {text}")
+                                    except Exception as e2:
+                                        result = {"message": f"鍐呴儴閿欒锛� {e2}", "type": "close"}
+                                        await websocket.send_json(result)
+                                        print(f"Error process message of ragflow: {e2}")
                 elif agent.type == "documentIa":
-                    print(122112)
                     token = DfTokenDao(db).get_token_by_id(DOCUMENT_IA_QUESTIONS)
                     # print(token)
                     if not token:

--
Gitblit v1.8.0