Newer
Older
cortex-hub / ai-hub / app / db / models / asset.py
from datetime import datetime
from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON, LargeBinary
from sqlalchemy.orm import relationship
from ..database import Base
import json
import base64
from hashlib import sha256
from cryptography.fernet import Fernet
from app.config import settings
from sqlalchemy.types import TypeDecorator

class EncryptedJSON(TypeDecorator):
    impl = Text

    def process_bind_param(self, value, dialect):
        if value is None:
            return None
        key = base64.urlsafe_b64encode(sha256(settings.SECRET_KEY.encode()).digest())
        fernet = Fernet(key)
        data = json.dumps(value).encode()
        encrypted = fernet.encrypt(data)
        return encrypted.decode('utf-8')

    def process_result_value(self, value, dialect):
        if value is None:
            return None
        key = base64.urlsafe_b64encode(sha256(settings.SECRET_KEY.encode()).digest())
        fernet = Fernet(key)
        decrypted = fernet.decrypt(value.encode('utf-8'))
        return json.loads(decrypted.decode('utf-8'))

class PromptTemplate(Base):
    __tablename__ = 'prompt_templates'

    id = Column(Integer, primary_key=True, index=True)
    slug = Column(String, unique=True, index=True, nullable=False)
    title = Column(String, nullable=False)
    content = Column(Text, nullable=False)
    version = Column(Integer, default=1)
    
    owner_id = Column(String, ForeignKey('users.id'), nullable=False, index=True)
    group_id = Column(String, ForeignKey('groups.id'), nullable=True)
    is_public = Column(Boolean, default=False)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    owner = relationship("User")
    group = relationship("Group")

    def __repr__(self):
        return f"<PromptTemplate(slug='{self.slug}', version={self.version})>"

class Skill(Base):
    __tablename__ = 'skills'

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, index=True, nullable=False)
    description = Column(String, nullable=True)
    skill_type = Column(String, default="local", nullable=False)
    
    is_enabled = Column(Boolean, default=True)
    features = Column(JSON, default=["chat"], nullable=True)
    
    owner_id = Column(String, ForeignKey('users.id'), nullable=False)
    is_system = Column(Boolean, default=False)
    
    extra_metadata = Column(JSON, default={}, nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    
    owner = relationship("User")
    files = relationship("SkillFile", back_populates="skill", cascade="all, delete-orphan")

    def __repr__(self):
        return f"<Skill(name='{self.name}', type='{self.skill_type}')>"

class SkillGroupAccess(Base):
    __tablename__ = 'skill_group_access'
    
    id = Column(Integer, primary_key=True, index=True)
    skill_id = Column(Integer, ForeignKey('skills.id'), nullable=False)
    group_id = Column(String, ForeignKey('groups.id'), nullable=False)
    
    granted_by = Column(String, ForeignKey('users.id'), nullable=False)
    granted_at = Column(DateTime, default=datetime.utcnow)
    
    skill = relationship("Skill")
    group = relationship("Group")

class SkillFile(Base):
    __tablename__ = 'skill_files'
    
    id = Column(Integer, primary_key=True, index=True)
    skill_id = Column(Integer, ForeignKey('skills.id'), nullable=False)
    file_path = Column(String, index=True, nullable=False)
    content = Column(Text, nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    skill = relationship("Skill", back_populates="files")

    def __repr__(self):
        return f"<SkillFile(skill_id={self.skill_id}, path='{self.file_path}')>"

class MCPServer(Base):
    __tablename__ = 'mcp_servers'

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    url = Column(String, nullable=False)
    auth_config = Column(EncryptedJSON, default={}, nullable=True)
    
    owner_id = Column(String, ForeignKey('users.id'), nullable=False)
    group_id = Column(String, ForeignKey('groups.id'), nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow)

    owner = relationship("User")
    group = relationship("Group")

    def __repr__(self):
        return f"<MCPServer(name='{self.name}', url='{self.url}')>"

class AssetPermission(Base):
    __tablename__ = 'asset_permissions'

    id = Column(Integer, primary_key=True, index=True)
    resource_type = Column(String, nullable=False, index=True)
    resource_id = Column(Integer, nullable=False, index=True)
    
    user_id = Column(String, ForeignKey('users.id'), nullable=True)
    group_id = Column(String, ForeignKey('groups.id'), nullable=True)
    
    access_level = Column(String, default="execute", nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)

    user = relationship("User")
    group = relationship("Group")

    def __repr__(self):
        return f"<AssetPermission(type='{self.resource_type}', id={self.resource_id}, level='{self.access_level}')>"