rag-gateway/app/models/group_model.py

from datetime import datetime
from enum import IntEnum
from typing import Optional

from sqlalchemy import Column, Integer, String, DateTime, Enum, Index, Table, ForeignKey
from pydantic import BaseModel
from sqlalchemy.orm import relationship, backref

from app.models.base_model import Base


group_knowledge_table = Table('group_Knowledge', Base.metadata
                           , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))
                           , Column('knowledge_id', String(32), ForeignKey('knowledgebase.id', ondelete='CASCADE')))

group_dialog_table = Table('group_dialogs', Base.metadata
                           , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))
                           , Column('dialog_id', String(36), ForeignKey('dialogs.id', ondelete='CASCADE')))

group_agent_table = Table('group_canvas', Base.metadata
                           , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))
                           , Column('canvas_id', String(36), ForeignKey('canvas.id', ondelete='CASCADE')))

group_llm_table = Table('group_llm', Base.metadata
                           , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))
                           , Column('llm_id', String(36), ForeignKey('common_llm.id', ondelete='CASCADE')))

group_unified_agent_table = Table('group_unified_agent', Base.metadata
                                    , Column('group_id', Integer, ForeignKey('group.id', ondelete='CASCADE'))
                                    ,Column('unified_agent_id', String(36), ForeignKey('unified_agent.id', ondelete='CASCADE')))


class GroupModel(Base):
    __tablename__ = "group"
    id = Column(Integer, primary_key=True, index=True)
    group_name = Column(String(255), unique=True, nullable=False, index=True)
    group_description = Column(String(255))
    status = Column(String(10), nullable=False, default="1")
    created_at = Column(DateTime, default=datetime.now())
    updated_at = Column(DateTime, default=datetime.now(), onupdate=datetime.now())
    creator = Column(Integer)
    group_type = Column(Integer, default=1)


    knowledges = relationship('KnowledgeModel',
                                 secondary=group_knowledge_table,
                                 backref=backref('groups', lazy='dynamic'),
                                 lazy="dynamic")

    dialogs = relationship('DialogModel',
                              secondary=group_dialog_table,
                              backref=backref('groups', lazy='dynamic'),
                              lazy="dynamic")

    agents = relationship('CanvasModel',
                             secondary=group_agent_table,
                             backref=backref('groups', lazy='dynamic'))

    llms = relationship('CommonLlmModel',
                           secondary=group_llm_table,
                           backref=backref('groups', lazy='dynamic'))

    unified_agents = relationship('UnifiedAgentModel',
                                     secondary=group_unified_agent_table,
                                     backref=backref('groups', lazy='dynamic'))
    def to_dict(self):
        return {
            'groupId': self.id,
            'name': self.group_name,
            'description': self.group_description,
            'status': self.status,
            'createTime': self.created_at.strftime("%Y-%m-%d %H:%M:%S")
        }


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


    def __repr__(self):
        return '<Role name:%r description:%r>\n' \
            % (self.group_name, self.group_description)

    def to_json(self):
        json = {
            'roleId': self.id,
            'createTime': self.created_at.strftime('%Y-%m-%d %H:%M:%S'),
            'updateTime': self.updated_at.strftime('%Y-%m-%d %H:%M:%S'),
            'groupName': self.group_name,
            'remark': self.group_description,
            'status': self.status,
        }
        if len(self.knowledges.all()) > 0:
            json['knowledges'] = [knowledge.to_json() for knowledge in self.knowledges]

        if len(self.dialogs.all()) > 0:
            json['dialogs'] = [dialog.to_json() for dialog in self.dialogs]

        if len(self.agents.all()) > 0:
            json['agents'] = [agent.to_json() for agent in self.agents]

        if len(self.llms.all()) > 0:
            json['llms'] = [llm.to_json() for llm in self.llms]

        json['users'] = [user.to_dict() for user in self.users]

        if hasattr(self, 'flag'):
            json['flag'] = self.flag

        return json


# class UserGroupModel(Base):
#     __tablename__ = "user_group"
#     id = Column(Integer, primary_key=True)
#     group_id = Column(Integer, nullable=False)
#     user_id = Column(Integer, nullable=False)
#     Index('ix_user_group_id', group_id, user_id, unique=True)


class GroupData(BaseModel):
    id: Optional[int] = None
    group_name: Optional[str] = ""
    group_description: Optional[str] = ""
    group_status: Optional[str] = ""

class GroupUsers(BaseModel):
    id: int
    user_list: list


class GroupResource(BaseModel):
    id: int
    agents: list
    dialogs: list
    knowledges: list