| | |
|---|
| | | 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 |
|---|
| | | DIFY = 4 |
|---|
| | | OTHER = 5 |
|---|
| | | |
|---|
| | | |
|---|
| | | 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 |
|---|
