Selaa lähdekoodia

请求basic_agent换成requests库,保存消息是不再保存中间结果

zhangqian 1 vuosi sitten
vanhempi
sitoutus
0ad20a6318
3 muutettua tiedostoa jossa 25 lisäystä ja 112 poistoa
  1. 13 4
      app/api/chat.py
  2. 0 69
      app/api/excel_talk.py
  3. 12 39
      app/service/basic.py

+ 13 - 4
app/api/chat.py

@@ -261,10 +261,19 @@ async def handle_client(websocket: WebSocket,
                                     f"&file_type={file_type}")
                         excel_url = build_file_url(excel_name, 'excel')
                         image_url = build_file_url(image_name, 'image')
-                        try:
-                            SessionService(db).update_session(chat_id, message={"content": output, "role": "assistant"})
-                        except Exception as e:
-                            logger.error(f"Unexpected error when update_session: {e}")
+                        if excel_url or data.get("e", ""):
+                            try:
+                                SessionService(db).update_session(chat_id,
+                                                                  message={
+                                                                      "content": output,
+                                                                      "excel_url": excel_url,
+                                                                      "image_url": image_url,
+                                                                      "sql": data.get("sql", ""),
+                                                                      "code": data.get("code", ""),
+                                                                      "e": data.get("e", ""),
+                                                                      "role": "assistant"})
+                            except Exception as e:
+                                logger.error(f"Unexpected error when update_session: {e}")
                         # 发送结果给客户端
                         data["type"] = "message"
                         data["message"] = output

+ 0 - 69
app/api/excel_talk.py

@@ -1,69 +0,0 @@
-import asyncio
-import json
-from enum import Enum
-
-from fastapi import APIRouter, Depends
-from sqlalchemy.orm import Session
-from starlette.websockets import WebSocket, WebSocketDisconnect
-
-from app.api import get_current_user_websocket
-from app.config.config import settings
-from app.models import UserModel, AgentModel
-from app.models.base_model import get_db
-from app.service.basic import BasicService
-
-router = APIRouter()
-
-# class CompletionRequest(BaseModel):
-#     id: Optional[str] = None
-#     app_id: str
-#     message: str
-#
-# class DownloadRequest(BaseModel):
-#     file_id: str
-#     app_id: str
-#     file_type: Optional[str] = None
-
-
-class AdvancedAgentID(Enum):
-    EXCEL_TALK = "excel_talk"
-    QUESTIONS_TALK = "questions_talk"
-
-@router.websocket("/ws/{agent_id}/{chat_id}")
-async def handle_client(websocket: WebSocket,
-                        agent_id: str,
-                        chat_id: str,
-                        current_user: UserModel = Depends(get_current_user_websocket),
-                        db: Session = Depends(get_db)):
-    await websocket.accept()
-    print(f"Client {agent_id} connected")
-
-    service = BasicService(base_url=settings.basic_base_url)
-
-    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
-    if not agent:
-        ret = {"message": "Agent not found", "type": "close"}
-        await websocket.send_json(ret)
-        return
-    try:
-        while True:
-            # 接收前端消息
-            message = await websocket.receive_json()
-            question = message.get("message")
-            if not question:
-                await websocket.send_json({"message": "Invalid request", "type": "error"})
-                continue
-
-            # 调用 excel_talk 方法
-            result = await service.excel_talk(question, chat_id)
-
-            # 将结果发送回前端
-            await websocket.send_json({"message": result, "type": "response"})
-    except Exception as e:
-        await websocket.send_json({"message": str(e), "type": "error"})
-    finally:
-        await websocket.close()
-        print(f"Client {agent_id} disconnected")
-
-
-

+ 12 - 39
app/service/basic.py

@@ -1,6 +1,7 @@
 import json
 
 import httpx
+import requests
 
 from Log import logger
 
@@ -48,8 +49,6 @@ class BasicService:
         url = f"{self.base_url}/exceltalk/upload/files"
         params = {'chat_id': chat_id, 'is_col': '0'}
 
-
-
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 url,
@@ -63,29 +62,26 @@ class BasicService:
         params = {'chat_id': chat_id}
         data = {"query": question}
         headers = {'Content-Type': 'application/json'}
-        buffer = bytearray()
-        async with httpx.AsyncClient(timeout=300.0) as client:
-            async with client.stream("POST", url, params=params, json=data, headers=headers) as response:
-                if response.status_code == 200:
+        with requests.post(url, headers=headers, json=data, params=params, timeout=60,
+                           stream=True) as response:
+            for line in response.iter_lines():
+                if line:
+                    decoded_line = line.decode("utf-8")
                     try:
-                        async for chunk in response.aiter_bytes():
-                            json_data = process_buffer(chunk, buffer)
-                            if json_data:
-                                yield json_data
-                                buffer.clear()
+                        if decoded_line.startswith("data:"):
+                            decoded_line = decoded_line[5:]
+                        answer = json.loads(decoded_line)
+                        yield answer
                     except GeneratorExit as e:
                         print(e)
                         yield {"message": "内部错误", "type": "close"}
                     finally:
-                        # 在所有数据接收完毕后记录日志
-                        logger.info("All messages received and processed - over")
+                        # 在所有数据接收完毕后返回close
                         yield {"message": "", "type": "close"}
 
                 else:
                     yield f"Error: {response.status_code}"
 
-
-
     async def questions_talk(self, question, chat_id: str):
         logger.error("---------------questions_talk--------------------------")
         url = f"{self.base_url}/questions/talk"
@@ -102,27 +98,4 @@ class BasicService:
 
     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={'word_name': file_id})
-
-
-def process_buffer(data, buffer):
-    def try_parse_json(data1):
-        try:
-            return True, json.loads(data1)
-        except json.JSONDecodeError:
-            return False, None
-
-    if data.startswith(b'data:'):
-        # 删除 'data:' 头
-        data = data[5:].strip()
-    else:
-        pass
-
-    # 直接拼接到缓冲区尝试解析JSON
-    buffer.extend(data.strip())
-    success, parsed_data = try_parse_json(buffer)
-    if success:
-        return parsed_data
-    else:
-        # 解析失败,继续拼接
-        return None
+        return await self.download_from_url(url, params={'word_name': file_id})