THOR4/asgard4/models.py

from app import db
from datetime import datetime
import json
from flask_login import UserMixin

class ThorClone(db.Model):
    """Model for Thor clones in the database"""
    id = db.Column(db.Integer, primary_key=True)
    clone_name = db.Column(db.String(100), unique=True, nullable=False)
    base_version = db.Column(db.String(100))
    description = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    modified_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    is_active = db.Column(db.Boolean, default=False)
    capabilities = db.Column(db.Text)  # JSON string of capabilities
    
    def get_capabilities(self):
        """Parse and return capabilities as a dictionary"""
        try:
            return json.loads(self.capabilities) if self.capabilities else {}
        except:
            return {}
    
    def set_capabilities(self, capabilities_dict):
        """Convert dictionary to JSON and store"""
        self.capabilities = json.dumps(capabilities_dict)
    
    def __repr__(self):
        return f"<ThorClone {self.clone_name}>"

class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(64), unique=True, nullable=False)
    email = db.Column(db.String(120), unique=True, nullable=False)
    password_hash = db.Column(db.String(256))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    # One-to-many relationship with conversations
    conversations = db.relationship('Conversation', backref='user', lazy='dynamic')
    
    # One-to-one relationship with settings
    settings = db.relationship('UserSettings', backref='user', uselist=False)
    
    def __repr__(self):
        return f'<User {self.username}>'

class Conversation(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
    title = db.Column(db.String(100), default="New Conversation")
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    # One-to-many relationship with messages
    messages = db.relationship('Message', backref='conversation', lazy='dynamic')
    
    def __repr__(self):
        return f'<Conversation {self.id} - {self.title}>'

class Message(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    conversation_id = db.Column(db.Integer, db.ForeignKey('conversation.id'))
    is_user = db.Column(db.Boolean, default=True)  # True if from user, False if from AI
    content = db.Column(db.Text, nullable=False)
    timestamp = db.Column(db.DateTime, default=datetime.utcnow)
    
    def __repr__(self):
        sender = "User" if self.is_user else "AI"
        return f'<Message {self.id} from {sender}>'

class UserSettings(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), unique=True)
    
    # AI safety and ethics settings (all default to True for safety)
    content_filter_enabled = db.Column(db.Boolean, default=True)
    ethics_check_enabled = db.Column(db.Boolean, default=True)
    permission_required = db.Column(db.Boolean, default=True)
    
    # Model preferences
    preferred_model = db.Column(db.String(100), default="gpt2")
    
    # Advanced settings stored as JSON
    advanced_settings = db.Column(db.Text, default='{}')
    
    def get_advanced_settings(self):
        """Parse and return advanced settings as a dictionary"""
        try:
            return json.loads(self.advanced_settings)
        except:
            return {}
    
    def set_advanced_settings(self, settings_dict):
        """Convert dictionary to JSON and store"""
        self.advanced_settings = json.dumps(settings_dict)
    
    def __repr__(self):
        return f'<UserSettings for user_id {self.user_id}>'

class ModelCache(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    model_name = db.Column(db.String(100), unique=True, nullable=False)
    source_url = db.Column(db.String(255))
    local_path = db.Column(db.String(255))
    downloaded_at = db.Column(db.DateTime, default=datetime.utcnow)
    last_used = db.Column(db.DateTime, default=datetime.utcnow)
    
    def __repr__(self):
        return f'<ModelCache {self.model_name}>'