Merge remote-tracking branch 'origin/master'

xuyonghao

2024-11-11 f4d4c9192087563156a92a32ed08edd54bf3b00a

 app/models/agent_model.py

@@ -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


@@ -25,3 +27,111 @@
            'agent_type': self.agent_type,
            'type': self.type
        }


class CanvasModel(Base):
    __tablename__ = 'user_canvas'
    __mapper_args__ = {
        # "order_by": 'SEQ'
    }
    id = Column(String(32), primary_key=True)
    create_time = Column(BigInteger)
    create_date = Column(DateTime, default=datetime.now)
    update_time = Column(BigInteger)
    update_date = Column(DateTime, default=datetime.now)
    avatar = Column(Text)
    user_id = Column(String(255))
    title = Column(String(255))
    description = Column(Text)
    canvas_type = Column(String(32))
    dsl = Column(Text)



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

    def to_json(self):
        return {
            'id': self.id,
            'create_time': self.create_time,
            'create_date': self.create_date,
            'update_time': self.update_time,
            '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 = UnifiedAgent.get_by_id(record_id)
        return record and record.type == t