Procházet zdrojové kódy

智能体列表增加排序和类型字段

zhangqian před 1 rokem
rodič
revize
cbb634f9c1
2 změnil soubory, kde provedl 19 přidání a 15 odebrání
  1. 2 10
      app/api/agent.py
  2. 17 5
      app/models/agent_model.py

+ 2 - 10
app/api/agent.py

@@ -26,15 +26,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)
 
 
@@ -64,4 +57,3 @@ async def chat_list(agent_id: str, db: Session = Depends(get_db), current_user:
 
     else:
         return ResponseList(code=200, msg="Unsupported agent type")
-

+ 17 - 5
app/models/agent_model.py

@@ -1,15 +1,27 @@
-from enum import IntEnum
-from sqlalchemy import Column, String, Enum as SQLAlchemyEnum
+from enum import IntEnum, StrEnum
+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
+        }