~xuyonghao/rag-gateway.git

			@@ -1,5 +1,7 @@
			import json
			from datetime import datetime
			from enum import IntEnum
			from sqlalchemy import Column, String, Enum as SQLAlchemyEnum, Integer
			from sqlalchemy import Column, String, Enum as SQLAlchemyEnum, Integer, BigInteger, DateTime, Text, Float, Boolean
			from app.models.base_model import Base


			@@ -7,6 +9,8 @@
			RAGFLOW = 1
			BISHENG = 2
			BASIC = 3
			DIFY = 4
			OTHER = 5


			class AgentModel(Base):
			@@ -25,3 +29,106 @@
			'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