出题组卷

zhaoqingang

2024-11-21 dbbc7d891c0f9837d3580842fa5326ba40a21476

 app/models/agent_model.py

@@ -1,15 +1,132 @@
import json
from datetime import datetime
from enum import IntEnum
from sqlalchemy import Column, String, Enum as SQLAlchemyEnum
from sqlalchemy import Column, String, Enum as SQLAlchemyEnum, Integer, BigInteger, DateTime, Text, Float, Boolean
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
        }


class CanvasModel(Base):
    __tablename__ = 'canvas'

    id = Column(String(32), primary_key=True)      # id
    create_date = Column(DateTime, default=datetime.now)
    update_date = Column(DateTime, default=datetime.now)
    avatar = Column(Text)                    # 图标
    user_id = Column(String(255))            # 用户id
    title = Column(String(255))               # 标题
    description = Column(Text)               # 说明
    canvas_type = Column(String(32))           # agent类型
    dsl = Column(Text)
    agent_type = Column(String(2))            # agent平台



    def get_id(self):
        return str(self.id)

    def to_json(self):
        return {
            'id': self.id,
            'create_date': self.create_date,
            'update_date': self.update_date,
            'avatar': self.avatar,
            'user_id': self.user_id,
            'title': self.title,
            'description': self.description,
            'canvas_type': self.canvas_type,
            'dsl': self.dsl
        }


class UnifiedAgentModel(Base):
    __tablename__ = 'unified_agent'

    id = Column(String(32), primary_key=True)
    tenant_id = Column(String(32))
    name = Column(String(255))
    description = Column(Text)
    icon = Column(Text)
    prompt_type = Column(String(16))
    prompt_config = Column(Text, default='{}')
    status = Column(String(1))
    prompts = Column(Text, default='[]')
    language = Column(String(32))
    llm_id = Column(String(128), default='')
    llm_setting = Column(Text, default='{}')
    similarity_threshold = Column(Float, default=0.0)
    vector_similarity_weight = Column(Float, default=0.0)
    top_n = Column(Integer, default=0)
    top_k = Column(Integer, default=0)
    do_refer = Column(String(1), default='')
    rerank_id = Column(String(128), default='')
    kb_ids = Column(Text, default='[]')
    hide = Column(Boolean, default=False)
    type = Column(Integer)
    canvas_type = Column(String(32), default="")
    dsl = Column(Text, default='{}')

    def get_id(self):
        return str(self.id)

    def to_json(self):
        if self.prompts is None or self.prompts == '':
            self.prompts = '[]'
        if self.prompt_config is None or self.prompt_config == '':
            self.prompt_config = '{}'
        if self.dsl is None or self.dsl == '':
            self.dsl = '{}'
        if self.kb_ids is None or self.kb_ids == '':
            self.kb_ids = '[]'
        return {
            'id': self.id,
            'create_time': self.create_time,
            'create_date': self.create_date,
            'update_time': self.update_time,
            'update_date': self.update_date,
            'tenant_id': self.tenant_id,
            'name': self.name,
            'description': self.description,
            'icon': self.icon,
            'language': self.language,
            'llm_id': self.llm_id,
            'similarity_threshold': self.similarity_threshold,
            'vector_similarity_weight': self.vector_similarity_weight,
            'top_n': self.top_n,
            'top_k': self.top_k,
            'do_refer': self.do_refer,
            'kb_ids': self.kb_ids,
            'status': self.status,
            'prompt_type': self.prompt_type,
            'prompt_config': json.loads(self.prompt_config),
            'prompts': json.loads(self.prompts),
            'dsl': json.loads(self.dsl)
        }

    @staticmethod
    def is_type(record_id, t):
        record = UnifiedAgentModel.get_by_id(record_id)
        return record and record.type == t