zhaoqingang
2024-12-03 319e3e93bbb4807da58db2ac72bb8f3029738dd7
app/api/chat.py
@@ -235,7 +235,7 @@
                        excel_url = None
                        if file_name:
                            excel_url = f"/api/files/download/?agent_id=basic_question_talk&file_id={file_name}&file_type=word"
                        result = {"message": output, "type": "message", "file_url": excel_url, "file_name":file_name}
                        result = {"message": output, "type": "message", "file_url": excel_url, "file_name": file_name}
                        try:
                            SessionService(db).update_session(chat_id,
                                                              message={"role": "assistant", "content": result})
@@ -268,6 +268,7 @@
                                return None
                            return (f"/api/files/download/?agent_id={agent_id}&file_id={name}"
                                    f"&file_type={file_type}")
                        if e_name:
                            excel_url = build_file_url(e_name, 'excel')
                            excel_name = e_name
@@ -294,7 +295,7 @@
                        await websocket.send_json(data)
                    if message_data:
                        try:
                            SessionService(db).update_session(chat_id,message=message_data)
                            SessionService(db).update_session(chat_id, message=message_data)
                        except Exception as e:
                            logger.error(f"Unexpected error when update_session: {e}")
        except Exception as e:
@@ -306,110 +307,214 @@
    if agent_type == AgentType.DIFY:
        dify_service = DifyService(settings.dify_base_url)
        # token = get_dify_token(db, current_user.id)
        token = settings.dify_api_token
        try:
            async def forward_to_dify():
                while True:
                    image_list = []
                    is_image = False
                    conversation_id = ""
                    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 = ""
                    answer_str = ""
                    async for rag_response in dify_service.chat(token, current_user.id, question, upload_file_id, conversation_id):
                        # print("=============================================")
                        # print(rag_response)
                if agent.type == "imageTalk":
                    token = settings.dify_api_token
                    while True:
                        image_list = []
                        is_image = False
                        conversation_id = ""
                        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:
                            if rag_response[:5] == "data:":
                                # 如果是,则截取掉前5个字符,并去除首尾空白符
                                complete_response = rag_response[5:].strip()
                            else:
                                # 否则,保持原样
                                complete_response = rag_response
                            # complete_response += text
                            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 = ""
                        answer_str = ""
                        async for rag_response in dify_service.chat(token, current_user.id, question, upload_file_id,
                                                                    conversation_id):
                            try:
                                data = json.loads(complete_response)
                                complete_response = ""
                                # data = json_data.get("data")
                                if data.get("event") == "agent_message":# "event": "message_end"
                                    if "answer" not in  data or not data["answer"]:  # 信息过滤
                                        logger.error("非法数据--------------------")
                                        # logger.error(data)
                                if rag_response[:5] == "data:":
                                    # 如果是,则截取掉前5个字符,并去除首尾空白符
                                    complete_response = rag_response[5:].strip()
                                else:
                                    # 否则,保持原样
                                    complete_response = rag_response
                                try:
                                    data = json.loads(complete_response)
                                    if data.get("event") == "agent_message":  # "event": "message_end"
                                        if "answer" not in data or not data["answer"]:  # 信息过滤
                                            logger.error("非法数据--------------------")
                                            # logger.error(data)
                                            continue
                                        else:  # 正常输出
                                            answer = data.get("answer", "")
                                            if isinstance(answer, str):
                                                if "![](https://res.stepfun.com/" in answer and image_list:
                                                    is_image = True
                                                    pattern = r'!\[\] *\(https://res\.stepfun\.com/image_gen/[^)]+\)'
                                                    url_image = image_list.pop()
                                                    new_answer = re.sub(pattern, url_image, answer)
                                                    answer_str += new_answer
                                                else:
                                                    answer_str += answer
                                            elif isinstance(answer, dict):
                                                logger.error("未知数据体:0---------------------------------")
                                                logger.error(answer)
                                                answer_str += answer.get("action_input", "")
                                            result = {"message": answer_str, "type": "message"}
                                    elif data.get("event") == "message_end":
                                        images_url = []
                                        if image_list and not is_image:
                                            answer_str += image_list[-1]
                                        result = {"message": answer_str,
                                                  "type": "close"}  # , "message_files": images_url
                                        try:
                                            SessionService(db).update_session(chat_id,
                                                                              message={"role": "assistant",
                                                                                       "content": {"answer": answer_str,
                                                                                                   "images": images_url}},
                                                                              conversation_id=data.get(
                                                                                  "conversation_id"))
                                        except Exception as e:
                                            logger.error("保存dify的会话异常!")
                                            logger.error(e)
                                    elif data.get("event") == "message_file":
                                        await  dify_service.save_images(data.get("url"), data.get("id") + ".png")
                                        image_list.append(f"![](/api/files/image/{data.get('id')})")
                                        # result = {"message": answer_str, "type": "message"}
                                        continue
                                    else:  # 正常输出
                                        answer = data.get("answer", "")
                                        if isinstance(answer, str):
                                            if "![](https://res.stepfun.com/" in answer and image_list:
                                                is_image = True
                                                pattern = r'!\[\] *\(https://res\.stepfun\.com/image_gen/[^)]+\)'
                                                url_image = image_list.pop()
                                                new_answer = re.sub(pattern, url_image, answer)
                                                answer_str += new_answer
                                            else:
                                                answer_str += answer
                                        elif isinstance(answer, dict):
                                            logger.error("未知数据体:0---------------------------------")
                                            logger.error(answer)
                                            answer_str += answer.get("action_input", "")
                                        result = {"message": answer_str, "type": "message"}
                                elif data.get("event") == "message_end":
                                    images_url = []
                                    # res_msg = await dify_service.get_session_history(token, data.get("conversation_id"), str(current_user.id))
                                    # if len(res_msg) > 0:
                                    #     message_files = res_msg[-1].get("message_files")
                                    #     for msg_file in message_files:
                                    #         await  dify_service.save_images(msg_file.get("url"), msg_file.get("id")+".png")
                                    #         images_url.append(msg_file.get("id"))
                                    # result = {"message": answer_str, "type": "close"} # , "message_files": images_url
                                    if image_list and not is_image:
                                        answer_str += image_list[-1]
                                    result = {"message": answer_str,
                                              "type": "close"}  # , "message_files": images_url
                                    try:
                                        SessionService(db).update_session(chat_id,
                                                                          message={"role": "assistant", "content": {"answer":answer_str, "images":images_url}},conversation_id=data.get("conversation_id"))
                                    except Exception as e:
                                        logger.error("保存dify的会话异常!")
                                        logger.error(e)
                                elif data.get("event") == "message_file":
                                    await  dify_service.save_images(data.get("url"), data.get("id") + ".png")
                                    image_list.append(f"![](/api/files/image/{data.get('id')})")
                                    # result = {"message": answer_str, "type": "message"}
                                    else:
                                        continue
                                    await websocket.send_json(result)
                                    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 == "reportWorkflow":
                    print(2323333232)
                    token = settings.dify_workflow_clean
                    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)
                        if not upload_files:
                            await websocket.send_json({"message": "Invalid request", "type": "error"})
                            continue
                        try:
                            session = SessionService(db).create_session(
                                chat_id,
                                title,
                                agent_id,
                                AgentType.DIFY,
                                current_user.id
                            )
                            conversation_id = session.conversation_id
                        except Exception as e:
                            logger.error(e)
                        inputs = {
                            "input_files": []
                        }
                        for file in upload_files:
                            inputs["input_files"].append({
                                "type": "document",
                                "transfer_method": "local_file",
                                "url": "",
                                "upload_file_id": file
                            })
                        if workflow_type == 2:
                            inputs["Completion_of_main_indicators"] = title
                            token = settings.dify_workflow_report
                        complete_response = ""
                        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:
                                    continue
                                await websocket.send_json(result)
                                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}")
                                    # 否则,保持原样
                                    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:  # 正常输出
                                            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", "")
                                            result = {"message": message, "type": "close", "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)
                                            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}")
            # 启动任务处理客户端消息
            tasks = [