From 992d91359f4e4437ddba9843173254441c896918 Mon Sep 17 00:00:00 2001
From: zhaoqingang <zhaoqg0118@163.com>
Date: 星期四, 02 一月 2025 18:03:28 +0800
Subject: [PATCH] 报告生成文档清洗

---
 app/api/chat.py |  345 +++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 files changed, 318 insertions(+), 27 deletions(-)

diff --git a/app/api/chat.py b/app/api/chat.py
index 3b14b8b..2993b2a 100644
--- a/app/api/chat.py
+++ b/app/api/chat.py
@@ -6,13 +6,17 @@
 import asyncio
 import websockets
 from sqlalchemy.orm import Session
+from starlette.responses import PlainTextResponse
 
 from Log import logger
 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_TO_REPORT_TITLE, \
+    DOCUMENT_TO_TITLE, DOCUMENT_IA_QUESTIONS
 from app.models.agent_model import AgentModel, AgentType
 from app.models.base_model import get_db
 from app.models.user_model import UserModel
+from app.service.common.api_token import DfTokenDao
 from app.service.dialog import update_session_history
 from app.service.basic import BasicService
 from app.service.difyService import DifyService
@@ -35,11 +39,14 @@
     print(f"Client {agent_id} connected")
 
     agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
+    print(agent_id)
     if not agent:
         ret = {"message": "Agent not found", "type": "close"}
         await websocket.send_json(ret)
         return
+
     agent_type = agent.agent_type
+    print(agent_type)
     if chat_id == "" or chat_id == "0":
         ret = {"message": "Chat ID not found", "type": "close"}
         await websocket.send_json(ret)
@@ -310,7 +317,10 @@
         try:
             async def forward_to_dify():
                 if agent.type == "imageTalk":
-                    token = settings.dify_api_token
+                    token = DfTokenDao(db).get_token_by_id(IMAGE_TO_TEXT)
+                    if not token:
+                        await websocket.send_json({"message": "Invalid token", "type": "error"})
+
                     while True:
                         image_list = []
                         is_image = False
@@ -334,9 +344,18 @@
                         except Exception as e:
                             logger.error(e)
                         # complete_response = ""
+                        files = []
+                        if upload_file_id:
+                            files.append({
+                                "type": "image",
+                                "transfer_method": "local_file",
+                                "url": "",
+                                "upload_file_id": upload_file_id
+                            })
                         answer_str = ""
-                        async for rag_response in dify_service.chat(token, current_user.id, question, upload_file_id,
-                                                                    conversation_id):
+                        async for rag_response in dify_service.chat(token, current_user.id, question, files,
+                                                                    conversation_id, {}):
+                            # print(rag_response)
                             try:
                                 if rag_response[:5] == "data:":
                                     # 濡傛灉鏄紝鍒欐埅鍙栨帀鍓�5涓瓧绗︼紝骞跺幓闄ら灏剧┖鐧界
@@ -403,21 +422,33 @@
                                 await websocket.send_json(result)
                                 print(f"Error process message of ragflow: {e2}")
                 elif agent.type == "reportWorkflow":
-                    print(2323333232)
-                    token = settings.dify_workflow_clean
+
+                    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}")
                         upload_files = receive_message.get('upload_files', [])
                         title = receive_message.get('title', "")
+                        sub_titles = receive_message.get('sub_titles', "")
                         workflow_type = receive_message.get('workflow', 1)
-                        if not upload_files:
-                            await websocket.send_json({"message": "Invalid request", "type": "error"})
-                            continue
+                        title_number = receive_message.get('title_number', 8)
+                        title_style = receive_message.get('title_style', "")
+                        title_query = receive_message.get('title_query', "")
+                        is_clean = receive_message.get('is_clean', 0)
+                        file_type = receive_message.get('file_type', 1)
+                        max_token = receive_message.get('max_tokens', 100000)
+                        tokens = receive_message.get('tokens', 0)
+                        if upload_files:
+                            title_query = "start"
+                        # if not upload_files:
+                            # await websocket.send_json({"message": "Invalid request", "type": "error"})
+                            # continue
                         try:
                             session = SessionService(db).create_session(
                                 chat_id,
-                                title,
+                                title if title else title_query,
                                 agent_id,
                                 AgentType.DIFY,
                                 current_user.id
@@ -429,21 +460,274 @@
                         }
                         files = []
                         for file in upload_files:
+                            if file_type == 1:
+                                files.append({
+                                    "type": "document",
+                                    "transfer_method": "local_file",
+                                    "url": "",
+                                    "upload_file_id": file
+                                })
+                            else:
+                                files.append({
+                                    "type": "document",
+                                    "transfer_method": "remote_url",
+                                    "url": file,
+                                    "upload_file_id": ""
+                                })
+                        inputs_list = []
+                        token_list = []
+                        if workflow_type == 1:
+                            inputs["input_files"] = files
+                            inputs_list.append(inputs)
+                            token_list.append(token)
+                        elif workflow_type == 2 and is_clean == 0:
+                            inputs["file_list"] = files
+                            inputs["Completion_of_main_indicators"] = title
+                            inputs["sub_titles"] = sub_titles
+                            token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_REPORT_TITLE)
+                            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["file_list"] = files
+                            inputs["number_of_title"] = title_number
+                            inputs["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_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)
+                            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)
+                            if not token:
+                                await websocket.send_json(
+                                    {"message": "Invalid token document_to_report", "type": "error"})
+                            inputs_list.append(inputs1)
+                            token_list.append(token)
+
+                        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(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
+                                        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父杈撳嚭
+                                                    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"}
+                                            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": ""}
+                                                elif isinstance(answer, dict):
+                                                    download_url = ""
+                                                    outputs = answer.get("outputs", {})
+                                                    if outputs:
+                                                        message = outputs.get("output", "")
+                                                        download_url = outputs.get("download_url", "")
+                                                    else:
+                                                        message = answer.get("error", "")
+                                                    if download_url:
+                                                        files = [{
+                                                            "type": "document",
+                                                            "transfer_method": "remote_url",
+                                                            "url": download_url,
+                                                            "upload_file_id": ""
+                                                        }]
+                                                    result = {"message": message, "type": "message", "download_url": download_url}
+                                                    try:
+                                                        SessionService(db).update_session(chat_id,
+                                                                                          message={"role": "assistant",
+                                                                                                   "content": {
+                                                                                                       "answer": message,
+                                                                                                       "download_url": download_url}},
+                                                                                          conversation_id=data.get(
+                                                                                              "conversation_id"))
+                                                    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("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
+                                                    result = {"message": "", "type": "close", "download_url": ""}
+
+
+                                            else:
+                                                continue
+                                            try:
+                                                await websocket.send_json(result)
+                                            except Exception  as e:
+                                                logger.error(e)
+                                                logger.error("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
+                                            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 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
+                                    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父杈撳嚭
+                                                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"}
+                                        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:
+                                                SessionService(db).update_session(chat_id,
+                                                                                  message={"role": "assistant",
+                                                                                           "content": {
+                                                                                               "answer": message,
+                                                                                               "download_url": ""}},
+                                                                                  conversation_id=data.get(
+                                                                                      "conversation_id"))
+                                            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}")
+                elif agent.type == "documentIa":
+                    print(122112)
+                    token = DfTokenDao(db).get_token_by_id(DOCUMENT_IA_QUESTIONS)
+                    # print(token)
+                    if not token:
+                        await websocket.send_json({"message": "Invalid token", "type": "error"})
+
+                    while True:
+                        conversation_id = ""
+                        # print(4343)
+                        receive_message = await websocket.receive_json()
+                        print(f"Received from client {chat_id}: {receive_message}")
+                        upload_file_id = receive_message.get('upload_file_id', [])
+                        question = receive_message.get('message', "")
+                        if not question and not image_url:
+                            await websocket.send_json({"message": "Invalid request", "type": "error"})
+                            continue
+                        try:
+                            session = SessionService(db).create_session(
+                                chat_id,
+                                question,
+                                agent_id,
+                                AgentType.DIFY,
+                                current_user.id
+                            )
+                            conversation_id = session.conversation_id
+                        except Exception as e:
+                            logger.error(e)
+                        # complete_response = ""
+                        files = []
+                        for fileId in upload_file_id:
                             files.append({
                                 "type": "document",
                                 "transfer_method": "local_file",
                                 "url": "",
-                                "upload_file_id": file
+                                "upload_file_id": fileId
                             })
-                        if workflow_type == 1:
-                            inputs["input_files"] = files
-                        if workflow_type == 2:
-                            inputs["file_list"] = files
-                            inputs["Completion_of_main_indicators"] = title
-                            token = settings.dify_workflow_report
+
+                        answer_str = ""
                         complete_response = ""
-                        async for rag_response in dify_service.workflow(token, current_user.id, inputs):
-                            print(rag_response)
+                        async for rag_response in dify_service.chat(token, current_user.id, question, files,
+                                                                    conversation_id, {}):
+                            # print(rag_response)
                             try:
                                 if rag_response[:5] == "data:":
                                     # 濡傛灉鏄紝鍒欐埅鍙栨帀鍓�5涓瓧绗︼紝骞跺幓闄ら灏剧┖鐧界
@@ -455,8 +739,8 @@
                                     complete_response += 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.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)
@@ -472,6 +756,11 @@
                                                 message = answer.get("title", "")
 
                                             result = {"message": message, "type": "system"}
+                                            continue
+                                    elif data.get("event") == "message":  # "event": "message_end"
+                                        # 姝e父杈撳嚭
+                                        answer = data.get("answer", "")
+                                        result = {"message": answer, "type": "stream"}
                                     elif data.get("event") == "workflow_finished":
                                         answer = data.get("data", "")
                                         if isinstance(answer, str):
@@ -482,12 +771,13 @@
                                             download_url = ""
                                             outputs = answer.get("outputs", {})
                                             if outputs:
-                                                message = outputs.get("output", "")
-                                                download_url = outputs.get("download_url", "")
+                                                message = outputs.get("answer", "")
+                                                # download_url = outputs.get("download_url", "")
                                             else:
                                                 message = answer.get("error", "")
 
-                                            result = {"message": message, "type": "message", "download_url": download_url}
+                                            # result = {"message": message, "type": "message",
+                                            #           "download_url": download_url}
                                             try:
                                                 SessionService(db).update_session(chat_id,
                                                                                   message={"role": "assistant",
@@ -499,15 +789,16 @@
                                             except Exception as e:
                                                 logger.error("淇濆瓨dify鐨勪細璇濆紓甯革紒")
                                                 logger.error(e)
-                                            await websocket.send_json(result)
-                                            result = {"message": "", "type": "close", "download_url": ""}
-
+                                            # await websocket.send_json(result)
+                                        continue
+                                    elif data.get("event") == "message_end":
+                                        result = {"message": "", "type": "close"}
 
                                     else:
                                         continue
                                     try:
                                         await websocket.send_json(result)
-                                    except Exception  as e:
+                                    except Exception as e:
                                         logger.error(e)
                                         logger.error("杩斿洖瀹㈡埛绔秷鎭紓甯�!")
                                     complete_response = ""

--
Gitblit v1.8.0