zhaoqingang 1 gadu atpakaļ
vecāks
revīzija
726a6682af
2 mainītis faili ar 26 papildinājumiem un 17 dzēšanām
  1. 6 5
      app/api/chat.py
  2. 20 12
      app/api/files.py

+ 6 - 5
app/api/chat.py

@@ -783,15 +783,15 @@ async def handle_client(websocket: WebSocket,
                             inputs["Question_Difficulty"] = receive_message["difficulty"]
                         if "is_paper" in receive_message:
                             inputs["Generate_test_paper"] = receive_message["is_paper"]
-                        if "single_choices" in receive_message:
+                        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_choices" in receive_message:
+                        if "multiple_choice" in receive_message:
                             inputs["Multiple_Choice"] = receive_message["multiple_choice"]
-                        if "easy_questions" in receive_message:
+                        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"]
@@ -823,6 +823,7 @@ async def handle_client(websocket: WebSocket,
                             })
                         if files:
                             inputs["upload_files"] = files
+                        # print(inputs)
                         if not question and not inputs:
                             await websocket.send_json({"message": "Invalid request", "type": "error"})
                             continue
@@ -844,7 +845,7 @@ async def handle_client(websocket: WebSocket,
                                     complete_response += rag_response
                                 try:
                                     data = json.loads(complete_response)
-                                    print(data)
+                                    # 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"]:  # 信息过滤
@@ -878,7 +879,7 @@ 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", "")
 

+ 20 - 12
app/api/files.py

@@ -10,7 +10,9 @@ from werkzeug.utils import send_file
 
 from app.api import Response, get_current_user, ResponseList
 from app.config.config import settings
-from app.config.const import DOCUMENT_TO_REPORT, IMAGE_TO_TEXT
+from app.config.const import DOCUMENT_TO_REPORT, IMAGE_TO_TEXT, DOCUMENT_TO_REPORT_TITLE, DOCUMENT_IA_QUESTIONS, \
+    DOCUMENT_TO_PAPER
+from app.models import MenuCapacityModel
 from app.models.agent_model import AgentType, AgentModel
 from app.models.base_model import get_db
 from app.models.user_model import UserModel
@@ -33,16 +35,17 @@ async def upload_files(
     db: Session = Depends(get_db),
     current_user: UserModel = Depends(get_current_user)
 ):
-    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
+    agent = db.query(MenuCapacityModel).filter(MenuCapacityModel.chat_id == agent_id).first()
     if not agent:
-        return Response(code=404, msg="Agent not found")
+        return ResponseList(code=404, msg="Agent not found")
+    agent_type = int(agent.capacity_type)
 
     # 检查 agent 类型,确定是否允许上传多个文件
-    if agent.agent_type in [AgentType.RAGFLOW, AgentType.BISHENG]:
+    if agent_type in [AgentType.RAGFLOW, AgentType.BISHENG]:
         if len(file) > 1:
             return Response(code=400, msg="这个智能体只支持传单个文件")
 
-    if agent.agent_type == AgentType.RAGFLOW or agent.agent_type == AgentType.BISHENG:
+    if agent_type == AgentType.RAGFLOW or agent_type == AgentType.BISHENG:
         file = file[0]
         # 读取上传的文件内容
         try:
@@ -50,7 +53,7 @@ async def upload_files(
         except Exception as e:
             return Response(code=400, msg=str(e))
 
-        if agent.agent_type == AgentType.RAGFLOW:
+        if agent_type == AgentType.RAGFLOW:
             token = await get_ragflow_token(db, current_user.id)
             ragflow_service = RagflowService(base_url=settings.fwr_base_url)
             # 查询会话是否存在,不存在先创建会话
@@ -63,7 +66,7 @@ async def upload_files(
             # 对于多文件,可能需要收集所有doc_ids
             return Response(code=200, msg="", data={"doc_ids": doc_ids, "file_name": file.filename})
 
-        elif agent.agent_type == AgentType.BISHENG:
+        elif agent_type == AgentType.BISHENG:
             bisheng_service = BishengService(base_url=settings.sgb_base_url)
             try:
                 token = await get_bisheng_token(db, current_user.id)
@@ -73,7 +76,7 @@ async def upload_files(
             result["file_name"] = file.filename
             return Response(code=200, msg="", data=result)
 
-    elif agent.agent_type == AgentType.BASIC:
+    elif agent_type == AgentType.BASIC:
         if agent_id == "basic_excel_talk":
             # 处理单个文件的情况
             file_list = file
@@ -97,9 +100,9 @@ async def upload_files(
             # service = BasicService(base_url=settings.basic_paper_url)
             # result = await service.paper_file_upload(chat_id, file.filename, file_content)
 
-    elif agent.agent_type == AgentType.DIFY:
+    elif agent_type == AgentType.DIFY:
         dify_service = DifyService(base_url=settings.dify_base_url)
-        if agent.type == "imageTalk":
+        if agent.chat_type == "imageTalk":
             token = DfTokenDao(db).get_token_by_id(IMAGE_TO_TEXT)
             if not token:
                 raise HTTPException(status_code=500, detail="获取token失败,image_and_text_conversion!")
@@ -113,8 +116,13 @@ async def upload_files(
                 data = await dify_service.upload(token, file.filename, file_content, current_user.id)
             except Exception as e:
                 raise HTTPException(status_code=500, detail=str(e))
-        elif agent.type == "reportWorkflow" or agent.type == "documentIa":
-            token = DfTokenDao(db).get_token_by_id(DOCUMENT_TO_REPORT)
+        elif agent.chat_type == "reportWorkflow" or agent.chat_type == "documentIa" or agent.chat_type == "paperTalk":
+            token_dict = {
+                "reportWorkflow": DOCUMENT_TO_REPORT,
+                "documentIa": DOCUMENT_IA_QUESTIONS,
+                "paperTalk": DOCUMENT_TO_PAPER,
+            }
+            token = DfTokenDao(db).get_token_by_id(token_dict[agent.chat_type])
             if not token:
                 raise HTTPException(status_code=500, detail="获取token失败,document_to_report!")
             result = []