智能体获取会话id接口

app/api/agent.py

@@ -1,3 +1,5 @@
+import uuid
+
 from fastapi import Depends, APIRouter, Query, HTTPException
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
@@ -57,3 +59,12 @@ async def chat_list(agent_id: str, db: Session = Depends(get_db), current_user:
 
     else:
         return ResponseList(code=200, msg="Unsupported agent type")
+
+
+@router.get("/get-chat-id/{agent_id}", response_model=Response)
+async def agent_list(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")
+
+    return Response(code=200, msg="", data={"chat_id": uuid.uuid4().hex})

app/api/chat.py

@@ -7,6 +7,7 @@ import websockets
 from sqlalchemy.orm import Session
 from app.api import get_current_user_websocket
 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.ragflow import RagflowService
@@ -14,9 +15,6 @@ from app.service.token import get_bisheng_token, get_ragflow_token
 
 router = APIRouter()
 
-# 存储客户端 WebSocket 连接
-client_websockets = {}
-
 
 # 中间层WebSocket 服务器，接收客户端的连接
 @router.websocket("/ws/{agent_id}/{chat_id}")
@@ -28,17 +26,16 @@ async def handle_client(websocket: WebSocket,
     await websocket.accept()
     print(f"Client {agent_id} connected")
 
-    if agent_id == "0":
-        agent_id = settings.bisheng_agent_id
-    elif agent_id == "1":
-        agent_id = settings.ragflow_agent_id
-        chat_id = settings.ragflow_chat_id
-
-    if chat_id == "0":
-        chat_id = uuid.uuid4().hex
+    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)
 
-    client_websockets[chat_id] = websocket
-    if agent_id == settings.ragflow_agent_id:
+    if agent_type == AgentType.RAGFLOW:
         ragflow_service = RagflowService(settings.ragflow_base_url)
         token = get_ragflow_token(db, current_user.id)
         try:
@@ -73,10 +70,8 @@ async def handle_client(websocket: WebSocket,
             await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
         except WebSocketDisconnect:
             print(f"Client {chat_id} disconnected")
-        finally:
-            del client_websockets[chat_id]
 
-    else:
+    elif agent_type == AgentType.BISHENG:
         token = get_bisheng_token(db, current_user.id)
         service_uri = f"{settings.bisheng_websocket_url}/api/v1/assistant/chat/{agent_id}?t=&chat_id={chat_id}"
         headers = {'cookie': f"access_token_cookie={token};"}
@@ -133,5 +128,7 @@ async def handle_client(websocket: WebSocket,
 
             except WebSocketDisconnect:
                 print(f"Client {chat_id} disconnected")
-            finally:
-                del client_websockets[chat_id]
+    else:
+        ret = {"message": "Agent not found", "type": "close"}
+        return websocket.send_json(ret)
+

app/service/ragflow.py

@@ -80,3 +80,22 @@ class RagflowService:
                 for item in data
             ]
             return result
+
+    async def set_session(self, token: str, dialog_id: str, name: str, chat_id: str, is_new: bool) -> bool:
+        url = f"{self.base_url}/v1/conversation/set?dialog_id={dialog_id}"
+        headers = {
+            "Authorization": token
+        }
+
+        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 False
+            return True
+

main.py

@@ -11,9 +11,9 @@ app = FastAPI(
   description="",
 )
 
-app.include_router(auth_router, prefix='/auth', tags=["auth"])
-app.include_router(chat_router, prefix='/chat', tags=["chat"])
-app.include_router(agent_router, prefix='/agent', tags=["agent"])
+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"])
 
 if __name__ == "__main__":
     import uvicorn