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

zhangqian

2024-10-15 cbb634f9c133c726a00faf8199ba8954e18d4b05

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

 app/api/agent.py

@@ -26,15 +26,8 @@

@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 @@

    else:
        return ResponseList(code=200, msg="Unsupported agent type")


 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
        }