Merge remote-tracking branch 'origin/master'

app/api/agent.py

@@ -16,16 +16,6 @@ from app.service.token import get_ragflow_token, get_bisheng_token
 router = APIRouter()
 
 
-# Pydantic 模型用于响应
-class AgentResponse(BaseModel):
-    id: str
-    name: str
-    agent_type: AgentType
-
-    class Config:
-        orm_mode = True
-
-
 @router.get("/list", response_model=ResponseList)
 async def agent_list(db: Session = Depends(get_db)):
     agents = db.query(AgentModel).order_by(AgentModel.sort.asc()).all()
@@ -62,7 +52,7 @@ async def chat_list(agent_id: str, db: Session = Depends(get_db), current_user:
 
 
 @router.get("/get-chat-id/{agent_id}", response_model=Response)
-async def agent_list(agent_id: str, db: Session = Depends(get_db)):
+async def get_chat_id(agent_id: str, db: Session = Depends(get_db)):
     agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
     if not agent:
         return Response(code=404, msg="Agent not found")

app/api/chat.py

@@ -29,11 +29,13 @@ async def handle_client(websocket: WebSocket,
     agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
     if not agent:
         ret = {"message": "Agent not found", "type": "close"}
-        return websocket.send_json(ret)
+        await websocket.send_json(ret)
+        return
     agent_type = agent.agent_type
     if chat_id == "" or chat_id == "0":
         ret = {"message": "Chat ID not found", "type": "close"}
-        return websocket.send_json(ret)
+        await websocket.send_json(ret)
+        return
 
     if agent_type == AgentType.RAGFLOW:
         ragflow_service = RagflowService(settings.ragflow_base_url)
@@ -45,31 +47,42 @@ async def handle_client(websocket: WebSocket,
                     print(f"Received from client {chat_id}: {message}")
                     chat_history = message.get('chatHistory', [])
                     if len(chat_history) == 0:
-
-                        chat_history = await ragflow_service.set_session(token, agent_id, message["message"], chat_id, True)
+                        chat_history = await ragflow_service.get_session_history(token, chat_id)
                         if len(chat_history) == 0:
-                            result = {"message": "内部错误：创建会话失败", "type": "close"}
-                            await websocket.send_json(result)
-                            continue
+                            chat_history = await ragflow_service.set_session(token, agent_id,
+                                                                             message["message"], chat_id, True)
+                            if len(chat_history) == 0:
+                                result = {"message": "内部错误：创建会话失败", "type": "close"}
+                                await websocket.send_json(result)
+                                await websocket.close()
+                                return
+                        else:
+                            chat_history.append({
+                                "content": message["message"],
+                                "role": "user"
+                            })
                     async for rag_response in ragflow_service.chat(token, chat_id, chat_history):
                         try:
                             print(f"Received from ragflow: {rag_response}")
                             if rag_response[:5] == "data:":
                                 # 如果是，则截取掉前5个字符，并去除首尾空白符
-                                json_str = rag_response[5:].strip()
+                                text = rag_response[5:].strip()
                             else:
                                 # 否则，保持原样
-                                json_str = rag_response
-                            json_data = json.loads(json_str)
-                            data = json_data.get("data")
-                            if data is True:  # 完成输出
-                                result = {"message": "", "type": "close"}
-                            elif data is None:  # 发生错误
-                                answer = json_data.get("retmsg", json_data.get("retcode"))
-                                result = {"message": "内部错误：" + answer, "type": "stream"}
-                            else:  # 正常输出
-                                answer = data.get("answer", "")
-                                result = {"message": answer, "type": "stream"}
+                                text = rag_response
+                            try:
+                                json_data = json.loads(text)
+                                data = json_data.get("data")
+                                if data is True:  # 完成输出
+                                    result = {"message": "", "type": "close"}
+                                elif data is None:  # 发生错误
+                                    answer = json_data.get("retmsg", json_data.get("retcode"))
+                                    result = {"message": "内部错误：" + answer, "type": "stream"}
+                                else:  # 正常输出
+                                    answer = data.get("answer", "")
+                                    result = {"message": answer, "type": "stream"}
+                            except json.JSONDecodeError:
+                                result = {"message": text, "type": "stream"}
                             await websocket.send_json(result)
                             print(f"Forwarded to client {chat_id}: {result}")
                         except Exception as e:
@@ -143,5 +156,5 @@ async def handle_client(websocket: WebSocket,
                 print(f"Client {chat_id} disconnected")
     else:
         ret = {"message": "Agent not found", "type": "close"}
-        return websocket.send_json(ret)
+        await websocket.send_json(ret)
 

app/api/files.py

@@ -0,0 +1,44 @@
+from fastapi import Depends, APIRouter, HTTPException, UploadFile, File, requests
+from sqlalchemy.orm import Session
+
+from app.api import Response, get_current_user, ResponseList
+from app.config.config import settings
+from app.models.agent_model import AgentType, AgentModel
+from app.models.base_model import get_db
+from app.models.user_model import UserModel
+from app.service.bisheng import BishengService
+from app.service.ragflow import RagflowService
+from app.service.token import get_ragflow_token, get_bisheng_token
+
+router = APIRouter()
+
+
+@router.post("/upload/{agent_id}", response_model=Response)
+async def upload_file(agent_id: str,
+                      file: UploadFile = File(...),
+                      db: Session = Depends(get_db),
+                      current_user: UserModel = Depends(get_current_user)
+                      ):
+    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
+    if not agent:
+        return Response(code=404, msg="Agent not found")
+    # 读取上传的文件内容
+    try:
+        file_content = await file.read()
+    except Exception as e:
+        return Response(code=400, msg=str(e))
+
+    if agent.agent_type == AgentType.RAGFLOW:
+        pass
+
+    elif agent.agent_type == AgentType.BISHENG:
+        bisheng_service = BishengService(base_url=settings.bisheng_base_url)
+        try:
+            token = get_bisheng_token(db, current_user.id)
+            result = await bisheng_service.upload(token, file.filename, file_content)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+        return Response(code=200, msg="", data=result)
+
+    else:
+        return Response(code=200, msg="Unsupported agent type")

app/api/report.py

@@ -0,0 +1,105 @@
+import json
+
+from fastapi import WebSocket, WebSocketDisconnect, APIRouter, Depends, HTTPException, Query
+import asyncio
+import websockets
+from sqlalchemy.orm import Session
+from app.api import get_current_user_websocket, ResponseList, get_current_user
+from app.config.config import settings
+from app.models.agent_model import AgentModel, AgentType
+from app.models.base_model import get_db
+from app.models.user_model import UserModel
+from app.service.bisheng import BishengService
+from app.service.token import get_bisheng_token
+
+router = APIRouter()
+
+
+@router.websocket("/ws/{agent_id}/{chat_id}")
+async def report_chat(websocket: WebSocket,
+                      agent_id: str,
+                      chat_id: str,
+                      current_user: UserModel = Depends(get_current_user_websocket),
+                      db: Session = Depends(get_db)):
+    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
+    if not agent:
+        ret = {"message": "Agent not found", "type": "close"}
+        return websocket.send_json(ret)
+    agent_type = agent.agent_type
+    if chat_id == "" or chat_id == "0":
+        ret = {"message": "Chat ID not found", "type": "close"}
+        return websocket.send_json(ret)
+
+    if agent_type != AgentType.BISHENG:
+        ret = {"message": "Agent error", "type": "close"}
+        return websocket.send_json(ret)
+
+    token = get_bisheng_token(db, current_user.id)
+    service_uri = f"{settings.bisheng_websocket_url}/api/v1/chat/{agent_id}?type=L1&t=&chat_id={chat_id}"
+    headers = {'cookie': f"access_token_cookie={token};"}
+
+    await websocket.accept()
+    print(f"Client {agent_id} connected")
+
+    async with websockets.connect(service_uri, extra_headers=headers) as service_websocket:
+
+        try:
+            # 处理客户端发来的消息
+            async def forward_to_service():
+                while True:
+                    message = await websocket.receive_json()
+                    print(f"Received from client, {chat_id}: {message}")
+                    # 添加 'agent_id' 和 'chat_id' 字段
+                    message['flow_id'] = agent_id
+                    message['chat_id'] = chat_id
+                    await service_websocket.send(json.dumps(message))
+                    print(f"Forwarded to bisheng: {message}")
+
+            # 监听毕昇发来的消息并转发给客户端
+            async def forward_to_client():
+                while True:
+                    message = await service_websocket.recv()
+                    print(f"Received from bisheng: {message}")
+                    data = json.loads(message)
+                    files = data.get("files", [])
+                    steps = data.get("intermediate_steps", "")
+                    if len(files) != 0 or steps != "" or data["type"] == "close":
+                        if data["type"] == "close":
+                            t = "close"
+                        else:
+                            t = "stream"
+                        result = {"step_message": steps, "type": t, "files": files}
+                        await websocket.send_json(result)
+                        print(f"Forwarded to client, {chat_id}: {result}")
+
+            # 启动两个任务，分别处理客户端和服务端的消息
+            tasks = [
+                asyncio.create_task(forward_to_service()),
+                asyncio.create_task(forward_to_client())
+            ]
+            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            # 取消未完成的任务
+            for task in pending:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+        except WebSocketDisconnect:
+            print(f"Client {chat_id} disconnected")
+
+
+@router.get("/variables/list", response_model=ResponseList)
+async def get_variables(agent_id: str = Query(..., description="The ID of the agent"), db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)):
+    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
+    if not agent:
+        return ResponseList(code=404, msg="Agent not found")
+    bisheng_service = BishengService(base_url=settings.bisheng_base_url)
+    try:
+        token = get_bisheng_token(db, current_user.id)
+        result = await bisheng_service.variable_list(token, agent_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    return ResponseList(code=200, msg="", data=result)

app/service/bisheng.py

@@ -1,5 +1,4 @@
 from datetime import datetime
-
 import httpx
 
 from app.config.config import settings
@@ -10,6 +9,21 @@ class BishengService:
     def __init__(self, base_url: str):
         self.base_url = base_url
 
+    def _check_response(self, response: httpx.Response):
+        if response.status_code not in [200, 201]:
+            raise Exception(f"Failed to fetch data from Bisheng API: {response.text}")
+        response_data = response.json()
+        status_code = response_data.get("status_code", 0)
+        if status_code != 200:
+            raise Exception(f"Failed to fetch data from Bisheng API: {response.text}")
+        # 检查返回的数据类型
+        if isinstance(response_data.get("data"), dict):
+            return response_data.get("data", {})
+        elif isinstance(response_data.get("data"), list):
+            return response_data.get("data", [])
+        else:
+            return {}
+
     async def register(self, username: str, password: str):
         public_key = await self.get_public_key_api()
         password = BishengCrypto(public_key, settings.PRIVATE_KEY).encrypt(password)
@@ -19,8 +33,7 @@ class BishengService:
                 json={"user_name": username, "password": password},
                 headers={'Content-Type': 'application/json'}
             )
-            if response.status_code != 200 and response.status_code != 201:
-                raise Exception(f"Bisheng registration failed: {response.text}")
+            self._check_response(response)
 
     async def login(self, username: str, password: str) -> str:
         public_key = await self.get_public_key_api()
@@ -31,9 +44,8 @@ class BishengService:
                 json={"user_name": username, "password": password},
                 headers={'Content-Type': 'application/json'}
             )
-            if response.status_code != 200 and response.status_code != 201:
-                raise Exception(f"Bisheng login failed: {response.text}")
-            return response.json().get('data', {}).get('access_token')
+            data = self._check_response(response)
+            return data.get('access_token')
 
     async def get_public_key_api(self) -> dict:
         async with httpx.AsyncClient() as client:
@@ -41,19 +53,16 @@ class BishengService:
                 f"{self.base_url}/api/v1/user/public_key",
                 headers={'Content-Type': 'application/json'}
             )
-            if response.status_code != 200:
-                raise Exception(f"Failed to get public key: {response.text}")
-            return response.json().get('data', {}).get('public_key')
+            data = self._check_response(response)
+            return data.get('public_key')
 
     async def get_chat_sessions(self, token: str) -> list:
         url = f"{self.base_url}/api/v1/chat/list?page=1&limit=40"
         headers = {'cookie': f"access_token_cookie={token};"}
         async with httpx.AsyncClient() as client:
             response = await client.get(url, headers=headers)
-            if response.status_code != 200:
-                raise Exception(f"Failed to fetch data from Bisheng API: {response.text}")
+            data = self._check_response(response)
 
-            data = response.json().get("data", [])
             result = [
                 {
                     "id": item["chat_id"],
@@ -63,3 +72,27 @@ class BishengService:
                 for item in data
             ]
             return result
+
+    async def variable_list(self, token: str, agent_id: str) -> list:
+        url = f"{self.base_url}/api/v1/variable/list?flow_id={agent_id}"
+        headers = {'cookie': f"access_token_cookie={token};"}
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, headers=headers)
+            data = self._check_response(response)
+            return data
+
+    async def upload(self, token: str, filename: str, file: bytes) -> dict:
+        url = f"{self.base_url}/api/v1/knowledge/upload"
+        headers = {'cookie': f"access_token_cookie={token};"}
+
+        # 创建表单数据，包含文件
+        files = {"file": (filename, file)}
+        async with httpx.AsyncClient() as client:
+            response = await client.post(url, headers=headers, files=files)
+            data = self._check_response(response)
+            file_path = data.get("file_path", "")
+            result = {
+                "file_path": file_path
+            }
+
+            return result

app/service/ragflow.py

@@ -1,5 +1,5 @@
 import httpx
-
+from typing import Union, Dict, List
 from app.config.config import settings
 from app.utils.rsa_crypto import RagflowCrypto
 
@@ -8,13 +8,30 @@ class RagflowService:
     def __init__(self, base_url: str):
         self.base_url = base_url
 
+    async def _handle_response(self, response: httpx.Response) -> Union[Dict, List]:
+        if response.status_code != 200:
+            return {}
+
+        data = response.json()
+        ret_code = data.get("retcode")
+        if ret_code != 0:
+            return {}
+
+        # 检查返回的数据类型
+        if isinstance(data.get("data"), dict):
+            return data.get("data", {})
+        elif isinstance(data.get("data"), list):
+            return data.get("data", [])
+        else:
+            return {}
+
     async def register(self, username: str, password: str):
         password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password)
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/v1/user/register",
-                json={"nickname": username, "email": f"{username}@example.com", "password": password},
-                headers={'Content-Type': 'application/json'}
+                headers={'Content-Type': 'application/json'},
+                json={"nickname": username, "email": f"{username}@example.com", "password": password}
             )
             if response.status_code != 200:
                 raise Exception(f"Ragflow registration failed: {response.text}")
@@ -24,12 +41,11 @@ class RagflowService:
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/v1/user/login",
-                json={"email": f"{username}@example.com", "password": password},
-                headers={'Content-Type': 'application/json'}
+                headers={'Content-Type': 'application/json'},
+                json={"email": f"{username}@example.com", "password": password}
             )
             if response.status_code != 200:
                 raise Exception(f"Ragflow login failed: {response.text}")
-                # 从响应头中提取 Authorization 字段
             authorization = response.headers.get('Authorization')
             if not authorization:
                 raise Exception("Authorization header not found in response")
@@ -40,17 +56,16 @@ class RagflowService:
             "conversation_id": chat_id,
             "messages": chat_history
         }
+
+        print(data)
         target_url = f"{self.base_url}/v1/conversation/completion"
-        async with httpx.AsyncClient() as client:
+        async with httpx.AsyncClient(timeout=300.0) as client:
             headers = {
                 'Content-Type': 'application/json',
                 'Authorization': token
             }
-            # 创建流式请求
             async with client.stream("POST", target_url, json=data, headers=headers) as response:
-                # 检查响应状态码
                 if response.status_code == 200:
-                    # 流式读取响应
                     try:
                         async for answer in response.aiter_text():
                             yield answer
@@ -62,15 +77,10 @@ class RagflowService:
 
     async def get_chat_sessions(self, token: str, dialog_id: str) -> list:
         url = f"{self.base_url}/v1/conversation/list?dialog_id={dialog_id}"
-        headers = {
-            "Authorization": token
-        }
+        headers = {"Authorization": token}
         async with httpx.AsyncClient() as client:
             response = await client.get(url, headers=headers)
-            if response.status_code != 200:
-                raise Exception(f"Failed to fetch data from Ragflow API: {response.text}")
-
-            data = response.json().get("data", [])
+            data = await self._handle_response(response)
             result = [
                 {
                     "id": item["id"],
@@ -83,26 +93,32 @@ class RagflowService:
 
     async def set_session(self, token: str, dialog_id: str, name: str, chat_id: str, is_new: bool) -> list:
         url = f"{self.base_url}/v1/conversation/set?dialog_id={dialog_id}"
-        headers = {
-            "Authorization": token
+        headers = {"Authorization": token}
+        data = {
+            "dialog_id": dialog_id,
+            "name": name,
+            "is_new": is_new,
+            "conversation_id": chat_id,
         }
-
-        data = {"dialog_id": dialog_id,
-                "name": name,
-                "is_new": is_new,
-                "conversation_id": chat_id,
-                }
-
         async with httpx.AsyncClient() as client:
             response = await client.post(url, headers=headers, json=data)
-            if response.status_code != 200:
-                return []
-            return [{
-                "content": "你好！ 我是你的助理，有什么可以帮到你的吗？",
-                "role": "assistant"
-            },
+            data = await self._handle_response(response)
+            return [
+                {
+                    "content": "你好！ 我是你的助理，有什么可以帮到你的吗？",
+                    "role": "assistant"
+                },
                 {
                     "content": name,
                     "doc_ids": [],
                     "role": "user"
-                }]
+                }
+            ] if data else []
+
+    async def get_session_history(self, token: str, chat_id: str) -> list:
+        url = f"{self.base_url}/v1/conversation/get?conversation_id={chat_id}"
+        headers = {"Authorization": token}
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, headers=headers)
+            data = await self._handle_response(response)
+            return data.get("message", [])

main.py

@@ -3,6 +3,8 @@ from app.api.auth import router as auth_router
 from app.api.chat import router as chat_router
 from app.api.agent import router as agent_router
 from app.api.excel import router as excel_router
+from app.api.files import router as files_router
+from app.api.report import router as report_router
 from app.models.base_model import init_db
 
 init_db()
@@ -16,6 +18,8 @@ app.include_router(auth_router, prefix='/api/auth', tags=["auth"])
 app.include_router(chat_router, prefix='/api/chat', tags=["chat"])
 app.include_router(agent_router, prefix='/api/agent', tags=["agent"])
 app.include_router(excel_router, prefix='/api/document', tags=["document"])
+app.include_router(files_router, prefix='/api/files', tags=["files"])
+app.include_router(report_router, prefix='/api/report', tags=["report"])
 
 if __name__ == "__main__":
     import uvicorn