from enum import IntEnum
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)
    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
        }