报告生成代码同步完善

app/api/chat.py

@@ -323,7 +323,7 @@ async def handle_client(websocket: WebSocket,
         # token = get_dify_token(db, current_user.id)
         try:
             async def forward_to_dify():
-                if chat_type == "imageTalk":
+                if agent.type == "imageTalk":
                     token = DfTokenDao(db).get_token_by_id(IMAGE_TO_TEXT)
                     if not token:
                         await websocket.send_json({"message": "Invalid token", "type": "error"})
@@ -351,8 +351,6 @@ async def handle_client(websocket: WebSocket,
                         except Exception as e:
                             logger.error(e)
                         # complete_response = ""
-
-                        answer_str = ""
                         files = []
                         if upload_file_id:
                             files.append({
@@ -361,6 +359,7 @@ async def handle_client(websocket: WebSocket,
                                 "url": "",
                                 "upload_file_id": upload_file_id
                             })
+                        answer_str = ""
                         async for rag_response in dify_service.chat(token, current_user.id, question, files,
                                                                     conversation_id, {}):
                             # print(rag_response)
@@ -429,18 +428,15 @@ async def handle_client(websocket: WebSocket,
                                 result = {"message": f"内部错误： {e2}", "type": "close"}
                                 await websocket.send_json(result)
                                 print(f"Error process message of ragflow: {e2}")
-                elif chat_type == "reportWorkflow":
+                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}")
                         upload_files = receive_message.get('upload_files', [])
                         title = receive_message.get('title', "")
-                        workflow_type = receive_message.get('workflow', 1)
                         sub_titles = receive_message.get('sub_titles', "")
+                        workflow_type = receive_message.get('workflow', 1)
                         title_number = receive_message.get('title_number', 8)
                         title_style = receive_message.get('title_style', "")
                         title_query = receive_message.get('title_query', "")
@@ -450,13 +446,19 @@ async def handle_client(websocket: WebSocket,
                         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
+                                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:
@@ -489,14 +491,14 @@ async def handle_client(websocket: WebSocket,
                             inputs["input_files"] = files
                             inputs["Completion_of_main_indicators"] = title
                             inputs_list.append({"inputs": inputs, "token": token, "workflow_type": workflow_type})
-                        if workflow_type == 2:
+                        elif workflow_type == 2:
                             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_cleaning", "type": "error"})
+                                    {"message": "Invalid token document_to_report", "type": "error"})
                             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
@@ -524,7 +526,8 @@ async def handle_client(websocket: WebSocket,
                                 await websocket.send_json(
                                     {"message": "Invalid token document_to_report", "type": "error"})
                             inputs_list.append({"inputs": inputs1, "token": token, "workflow_type": 3})
-                        complete_response = ""
+
+                        # print(inputs_list)
                         for idx, input in enumerate(inputs_list):
                             # print(input)
                             if idx < len(inputs_list) - 1:
@@ -803,8 +806,7 @@ async def handle_client(websocket: WebSocket,
                                         result = {"message": f"内部错误： {e2}", "type": "close"}
                                         await websocket.send_json(result)
                                         print(f"Error process message of ragflow: {e2}")
-                elif chat_type == "documentIa":
-                    # print(122112)
+                elif agent.type == "documentIa":
                     token = DfTokenDao(db).get_token_by_id(DOCUMENT_IA_QUESTIONS)
                     # print(token)
                     if not token:
@@ -845,162 +847,6 @@ async def handle_client(websocket: WebSocket,
                         complete_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个字符，并去除首尾空白符
-                                    complete_response = rag_response[5:].strip()
-                                elif "event: ping" in rag_response:
-                                    continue
-                                else:
-                                    # 否则，保持原样
-                                    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"
-                                        if "data" not in data or not data["data"]:  # 信息过滤
-                                            logger.error("非法数据--------------------")
-                                            logger.error(data)
-                                            continue
-                                        else:  # 正常输出
-                                            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"}
-                                            # continue
-                                    elif data.get("event") == "message":  # "event": "message_end"
-                                         # 正常输出
-                                        answer = data.get("answer", "")
-                                        result = {"message": answer, "type": "stream"}
-                                    elif data.get("event") == "error":
-                                        answer = data.get("message", "")
-                                        result = {"message": answer, "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("answer", "")
-                                                # download_url = outputs.get("download_url", "")
-                                            else:
-                                                message = answer.get("error", "")
-
-                                            result = {"message": message, "type": "system",
-                                                      "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)
-                                            # 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:
-                                        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 chat_type == "paperTalk":
-                    token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_PAPER)
-                    # print(token)
-                    if not token:
-                        await websocket.send_json({"message": "Invalid token", "type": "error"})
-
-                    while True:
-                        conversation_id = ""
-                        inputs = {}
-                        # print(4343)
-                        receive_message = await websocket.receive_json()
-                        print(f"Received from client {chat_id}: {receive_message}")
-                        if "difficulty" in receive_message:
-                            inputs["Question_Difficulty"] = receive_message["difficulty"]
-                        if "is_paper" in receive_message:
-                            inputs["Generate_test_paper"] = receive_message["is_paper"]
-                        if "single_choice" in receive_message:
-                            inputs["Multiple_choice_questions"] = receive_message["single_choice"]
-                        if "gap_filling" in receive_message:
-                            inputs["Fill_in_blank"] = receive_message["gap_filling"]
-                        if "true_or_false" in receive_message:
-                            inputs["true_or_false"] = receive_message["true_or_false"]
-                        if "multiple_choice" in receive_message:
-                            inputs["Multiple_Choice"] = receive_message["multiple_choice"]
-                        if "easy_question" in receive_message:
-                            inputs["Short_Answer_Questions"] = receive_message["easy_question"]
-                        if "case_questions" in receive_message:
-                            inputs["Case_Questions"] = receive_message["case_questions"]
-                        if "key_words" in receive_message:
-                            inputs["key_words"] = receive_message["key_words"]
-                        upload_files = receive_message.get('upload_files', [])
-                        question = receive_message.get('message', "")
-                        session_log = SessionService(db).get_session_by_id(chat_id)
-                        if not session_log and not upload_files:
-                            await websocket.send_json({"message": "需要上传文档！", "type": "error"})
-                            continue
-                        try:
-                            session = SessionService(db).create_session(
-                                chat_id,
-                                question if question else "开始出题",
-                                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_files:
-                            files.append({
-                                "type": "document",
-                                "transfer_method": "local_file",
-                                "url": "",
-                                "upload_file_id": fileId
-                            })
-                        if files:
-                            inputs["upload_files"] = files
-                        # print(inputs)
-                        if not question and not inputs:
-                            await websocket.send_json({"message": "Invalid request", "type": "error"})
-                            continue
-
-                        if not question:
-                            question = "开始出题"
-                        complete_response = ""
-                        async for rag_response in dify_service.chat(token, current_user.id, question, files,
-                                                                    conversation_id, inputs):
                             # print(rag_response)
                             try:
                                 if rag_response[:5] == "data:":
@@ -1013,7 +859,6 @@ async def handle_client(websocket: WebSocket,
                                     complete_response += rag_response
                                 try:
                                     data = json.loads(complete_response)
-                                    # print(data)
                                     if data.get("event") == "node_started" or data.get(
                                             "event") == "node_finished":  # "event": "message_end"
                                         if "data" not in data or not data["data"]:  # 信息过滤
@@ -1031,14 +876,11 @@ async def handle_client(websocket: WebSocket,
                                                 message = answer.get("title", "")
 
                                             result = {"message": message, "type": "system"}
-                                            # continue
+                                            continue
                                     elif data.get("event") == "message":  # "event": "message_end"
-                                         # 正常输出
+                                        # 正常输出
                                         answer = data.get("answer", "")
                                         result = {"message": answer, "type": "stream"}
-                                    elif data.get("event") == "error":
-                                        answer = data.get("message", "")
-                                        result = {"message": answer, "type": "system"}
                                     elif data.get("event") == "workflow_finished":
                                         answer = data.get("data", "")
                                         if isinstance(answer, str):
@@ -1050,12 +892,12 @@ async def handle_client(websocket: WebSocket,
                                             outputs = answer.get("outputs", {})
                                             if outputs:
                                                 message = outputs.get("answer", "")
-                                                download_url = outputs.get("download_url", "")
+                                                # download_url = outputs.get("download_url", "")
                                             else:
                                                 message = answer.get("error", "")
 
-                                            result = {"message": message, "type": "system",
-                                                      "download_url": download_url}
+                                            # result = {"message": message, "type": "message",
+                                            #           "download_url": download_url}
                                             try:
                                                 SessionService(db).update_session(chat_id,
                                                                                   message={"role": "assistant",
@@ -1068,7 +910,7 @@ async def handle_client(websocket: WebSocket,
                                                 logger.error("保存dify的会话异常！")
                                                 logger.error(e)
                                             # await websocket.send_json(result)
-                                        # continue
+                                        continue
                                     elif data.get("event") == "message_end":
                                         result = {"message": "", "type": "close"}
 

app/models/session_model.py

@@ -24,6 +24,7 @@ class SessionModel(Base):
     tenant_id = Column(Integer)  # 创建人
     message = Column(TEXT)  # 说明
     conversation_id = Column(String(64))
+    workflow = Column(Integer, default=0)
 
     # to_dict 方法
     def to_dict(self):

app/service/session.py

@@ -12,8 +12,8 @@ class SessionService:
     def __init__(self, db: Session):
         self.db = db
 
-    def create_session(self, session_id: str, name: str, agent_id: str, agent_type: AgentType, user_id: int) -> Type[
-                                                                                                                    SessionModel] | SessionModel:
+    def create_session(self, session_id: str, name: str, agent_id: str, agent_type: AgentType, user_id: int,
+                       message: dict = None, workflow_type: int = 0) -> Type[SessionModel] | SessionModel:
         """
         创建一个新的会话记录。
 
@@ -26,9 +26,11 @@ class SessionService:
         返回:
             SessionModel: 新创建的会话模型实例，如果会话ID已存在则返回None。
         """
+        if not message:
+            message = {"role": "user", "content": name}
         existing_session = self.get_session_by_id(session_id)
         if existing_session:
-            existing_session.add_message({"role": "user", "content": name})
+            existing_session.add_message(message)
             existing_session.update_date = current_time()
             self.db.commit()
             self.db.refresh(existing_session)
@@ -40,7 +42,8 @@ class SessionService:
             agent_id=agent_id,
             agent_type=agent_type,
             tenant_id=user_id,
-            message=json.dumps([{"role": "user", "content": name}])
+            workflow=workflow_type,
+            message=json.dumps([message])
         )
         self.db.add(new_session)
         self.db.commit()