zhaoqingang 1 年間 前
コミット
dbbc7d891c

+ 54 - 28
app/api/chat.py

@@ -200,6 +200,7 @@ async def handle_client(websocket: WebSocket,
                             pass
     elif agent_type == AgentType.BASIC:
         try:
+            service = BasicService(base_url=settings.basic_base_url)
             while True:
                 # 接收前端消息
                 message = await websocket.receive_json()
@@ -213,36 +214,61 @@ async def handle_client(websocket: WebSocket,
                 if not question:
                     await websocket.send_json({"message": "Invalid request", "type": "error"})
                     continue
+                if agent.agent_type == "questionTalk":
+                    async for result in service.questions_talk(question, chat_id):
+                        try:
+                            if result[:5] == "data:":
+                                # 如果是,则截取掉前5个字符,并去除首尾空白符
+                                text = result[5:].strip()
+                            else:
+                                # 否则,保持原样
+                                text = result
+                            try:
+                                data = json.loads(text)
+                                output = data.get("output", "")
+                                file_name = data.get("filename", "")
 
-                service = BasicService(base_url=settings.basic_base_url)
-                async for result in service.excel_talk(question, chat_id):
-                    try:
-                        if result[:5] == "data:":
-                            # 如果是,则截取掉前5个字符,并去除首尾空白符
-                            text = result[5:].strip()
-                        else:
-                            # 否则,保持原样
-                            text = result
+                                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}
+                                await websocket.send_json(result | data)
+                            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 basic agent: {e2}")
+                else:
+                    async for result in service.excel_talk(question, chat_id):
                         try:
-                            data = json.loads(text)
-                            output = data.get("output", "")
-                            excel_name = data.get("excel_name", "")
-                            image_name = data.get("image_name", "")
-                            excel_url = None
-                            image_url = None
-                            if excel_name:
-                                excel_url = f"/api/files/download/?agent_id=basic_excel_talk&file_id={excel_name}&file_type=excel"
-                            if image_name:
-                                image_url = f"/api/files/download/?agent_id=basic_excel_talk&file_id={image_name}&file_type=image"
-                            result = {"message": output, "type": "message", "excel_url": excel_url, "image_url": image_url}
-                            await websocket.send_json(result | data)
-                        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 basic agent: {e2}")
+                            if result[:5] == "data:":
+                                # 如果是,则截取掉前5个字符,并去除首尾空白符
+                                text = result[5:].strip()
+                            else:
+                                # 否则,保持原样
+                                text = result
+                            try:
+                                data = json.loads(text)
+                                output = data.get("output", "")
+                                excel_name = data.get("excel_name", "")
+                                image_name = data.get("image_name", "")
+                                excel_url = None
+                                image_url = None
+                                if excel_name:
+                                    excel_url = f"/api/files/download/?agent_id=basic_excel_talk&file_id={excel_name}&file_type=excel"
+                                if image_name:
+                                    image_url = f"/api/files/download/?agent_id=basic_excel_talk&file_id={image_name}&file_type=image"
+                                result = {"message": output, "type": "message", "excel_url": excel_url, "image_url": image_url}
+                                await websocket.send_json(result | data)
+                            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 basic agent: {e2}")
         except Exception as e:
             await websocket.send_json({"message": str(e), "type": "error"})
         finally:

+ 15 - 0
app/api/files.py

@@ -68,6 +68,12 @@ async def upload_file(agent_id: str,
 
             return Response(code=200, msg="", data=result)
 
+        elif agent_id == "basic_paper_agent":
+            service = BasicService(base_url=settings.basic_paper_url)
+            result = await service.paper_file_upload(chat_id, file.filename, file_content)
+
+            return Response(code=200, msg="", data=result)
+
     else:
         return Response(code=200, msg="Unsupported agent type")
 
@@ -101,6 +107,8 @@ async def download_file(
         if agent_id == "basic_excel_talk":
             return await download_basic_file(file_id, file_type)
 
+        elif agent_id == "basic_question_talk":
+            return await download_basic_file(file_id, file_type)
     else:
         return Response(code=400, msg="Unsupported agent type")
 
@@ -137,5 +145,12 @@ async def download_basic_file(file_id: str, file_type: str):
             media_type=mimetype,
             headers={"Content-Disposition": f"attachment; filename={filename}"}
         )
+    elif file_type == "word":
+        content, filename, mimetype = await service.questions_talk_word_download(file_id)
+        return StreamingResponse(
+            io.BytesIO(content),
+            media_type=mimetype,
+            headers={"Content-Disposition": f"attachment; filename={filename}"}
+        )
     else:
         return Response(code=400, msg="Unsupported file type")

+ 1 - 0
app/config/config.py

@@ -16,6 +16,7 @@ class Settings:
     PRIVATE_KEY: str
     PASSWORD_KEY: str
     basic_base_url: str = ''
+    basic_paper_url: str = ''
     def __init__(self, **kwargs):
         # Check if all required fields are provided and set them
         for field in self.__annotations__.keys():

+ 2 - 1
app/config/config.yaml

@@ -2,7 +2,7 @@ secret_key: your-secret-key
 sgb_base_url: http://192.168.20.119:13001
 sgb_websocket_url: ws://192.168.20.119:13001
 fwr_base_url: http://192.168.20.119:11080
-database_url: mysql+pymysql://root:infini_rag_flow@192.168.20.116:5455/rag_basic
+database_url: mysql+pymysql://root:infini_rag_flow@192.168.20.119:5455/rag_basic
 sgb_db_url: mysql+pymysql://root:1234@192.168.20.119:13306/bisheng
 fwr_db_url: mysql+pymysql://root:infini_rag_flow@192.168.20.119:15455/rag_flow
 PUBLIC_KEY: |
@@ -14,3 +14,4 @@ fetch_sgb_agent: 报告生成,文档智能
 fetch_fwr_agent: 知识问答,智能问答
 PASSWORD_KEY: VKinqB-8XMrwCLLrcf_PyHyo12_4PVKvWzaHjNFions=
 basic_base_url: http://192.168.20.231:8000
+basic_paper_url: http://192.168.20.231:8000

+ 1 - 1
app/models/agent_model.py

@@ -128,5 +128,5 @@ class UnifiedAgentModel(Base):
 
     @staticmethod
     def is_type(record_id, t):
-        record = UnifiedAgent.get_by_id(record_id)
+        record = UnifiedAgentModel.get_by_id(record_id)
         return record and record.type == t

+ 17 - 1
app/service/basic.py

@@ -71,4 +71,20 @@ class BasicService:
                         print(e)
                         return
                 else:
-                    yield f"Error: {response.status_code}"
+                    yield f"Error: {response.status_code}"
+
+    async def questions_talk(self, chat_id: str):
+        url = f"{self.base_url}/questions/talk"
+        params = {'chat_id': chat_id}
+        headers = {'Content-Type': 'text/plain'}
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                url,
+                headers=headers,
+                params=params
+            )
+            return self._check_response(response)
+
+    async def questions_talk_word_download(self, file_id: str):
+        url = f"{self.base_url}/questions/download/word"
+        return await self.download_from_url(url, params={'excel_name': file_id})

+ 2 - 1
app/task/fetch_agent.py

@@ -115,7 +115,8 @@ def initialize_agents():
             ('bfd090d589d811efb3630242ac190006', 4, '文档智能', 'BISHENG', 'report'),
             ('da3451da89d911efb9490242ac190006', 3, '知识问答', 'RAGFLOW', 'knowledgeQA'),
             ('e96eb7a589db11ef87d20242ac190006', 5, '智能问答', 'RAGFLOW', 'chat'),
-            ('basic_excel_talk', 6, '智能数据', 'BASIC', 'excelTalk')
+            ('basic_excel_talk', 6, '智能数据', 'BASIC', 'excelTalk'),
+            ('basic_question_talk', 7, '文档出题', 'BASIC', 'questionTalk')
         ]
 
         for agent in initial_agents: