Merge remote-tracking branch 'origin/master'

# Conflicts:
#	main.py

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
@@ -26,15 +28,8 @@ class AgentResponse(BaseModel):
 
 @router.get("/list", response_model=ResponseList)
 async def agent_list(db: Session = Depends(get_db)):
-    agents = db.query(AgentModel).all()
-    result = [
-        {
-            "id": item.id,
-            "name": item.name,
-            "agent_type": item.agent_type
-        }
-        for item in agents
-    ]
+    agents = db.query(AgentModel).order_by(AgentModel.sort.asc()).all()
+    result = [item.to_dict() for item in agents]
     return ResponseList(code=200, msg="", data=result)
 
 
@@ -65,3 +60,11 @@ 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/models/agent_model.py

@@ -1,15 +1,27 @@
 from enum import IntEnum
-from sqlalchemy import Column, String, Enum as SQLAlchemyEnum
+from sqlalchemy import Column, String, Enum as SQLAlchemyEnum, Integer
 from app.models.base_model import Base
 
 
 class AgentType(IntEnum):
     RAGFLOW = 1
     BISHENG = 2
+    BASIC = 3
 
 
 class AgentModel(Base):
     __tablename__ = "agent"
-    id = Column(String(255), primary_key=True, index=True)
-    name = Column(String(255), index=True)
-    agent_type = Column(SQLAlchemyEnum(AgentType), nullable=False)  # 1 ragflow 2 bisheng
+    id = Column(String(255), primary_key=True)
+    name = Column(String(255))
+    sort = Column(Integer, default=0, nullable=False)
+    agent_type = Column(SQLAlchemyEnum(AgentType), nullable=False)
+    type = Column(String(255), nullable=False)
+
+    # to_dict 方法
+    def to_dict(self):
+        return {
+            'id': self.id,
+            'name': self.name,
+            'agent_type': self.agent_type,
+            'type': self.type
+        }

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

@@ -12,9 +12,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"])
 app.include_router(excel_router, prefix='/document', tags=["document"])
 
 if __name__ == "__main__":