from enum import IntEnum
from sqlalchemy import Column, String, Enum as SQLAlchemyEnum
from app.models.base_model import Base


class AgentType(IntEnum):
    RAGFLOW = 1
    BISHENG = 2


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