eveAI/common/models/user.py

from common.extensions import db
from flask_security import UserMixin, RoleMixin
from sqlalchemy.dialects.postgresql import ARRAY
import sqlalchemy as sa
from sqlalchemy import CheckConstraint


class Tenant(db.Model):
    """Tenant model"""

    __bind_key__ = 'public'
    __table_args__ = {'schema': 'public'}

    # Versioning Information
    created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())

    # company Information
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(80), unique=True, nullable=False)
    website = db.Column(db.String(255), nullable=True)

    # language information
    default_language = db.Column(db.String(2), nullable=True)
    allowed_languages = db.Column(ARRAY(sa.String(2)), nullable=True)

    # LLM specific choices
    embedding_model = db.Column(db.String(50), nullable=True)
    llm_model = db.Column(db.String(50), nullable=True)

    # Embedding variables
    html_tags = db.Column(ARRAY(sa.String(10)), nullable=True, default=['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'li'])
    html_end_tags = db.Column(ARRAY(sa.String(10)), nullable=True, default=['p', 'li'])
    html_included_elements = db.Column(ARRAY(sa.String(50)), nullable=True)
    html_excluded_elements = db.Column(ARRAY(sa.String(50)), nullable=True)

    # Embedding search variables
    es_k = db.Column(db.Integer, nullable=True, default=5)
    es_similarity_threshold = db.Column(db.Float, nullable=True, default=0.7)

    # Chat variables
    chat_RAG_temperature = db.Column(db.Float, nullable=True, default=0.3)
    chat_no_RAG_temperature = db.Column(db.Float, nullable=True, default=0.5)

    # Licensing Information
    license_start_date = db.Column(db.Date, nullable=True)
    license_end_date = db.Column(db.Date, nullable=True)
    allowed_monthly_interactions = db.Column(db.Integer, nullable=True)
    encrypted_chat_api_key = db.Column(db.String(500), nullable=True)

    # Tuning enablers
    embed_tuning = db.Column(db.Boolean, nullable=True, default=False)
    rag_tuning = db.Column(db.Boolean, nullable=True, default=False)

    # Relations
    users = db.relationship('User', backref='tenant')
    domains = db.relationship('TenantDomain', backref='tenant')

    def __repr__(self):
        return f"<Tenant {self.id}: {self.name}>"

    def to_dict(self):
        return {
            'id': self.id,
            'name': self.name,
            'website': self.website,
            'default_language': self.default_language,
            'allowed_languages': self.allowed_languages,
            'embedding_model': self.embedding_model,
            'llm_model': self.llm_model,
            'license_start_date': self.license_start_date,
            'license_end_date': self.license_end_date,
            'allowed_monthly_interactions': self.allowed_monthly_interactions
        }


class Role(db.Model, RoleMixin):
    __bind_key__ = 'public'
    __table_args__ = {'schema': 'public'}

    id = db.Column(db.Integer(), primary_key=True)
    name = db.Column(db.String(80), unique=True)
    description = db.Column(db.String(255))


class RolesUsers(db.Model):
    __bind_key__ = 'public'
    __table_args__ = {'schema': 'public'}

    user_id = db.Column(db.Integer(), db.ForeignKey('public.user.id', ondelete='CASCADE'), primary_key=True)
    role_id = db.Column(db.Integer(), db.ForeignKey('public.role.id', ondelete='CASCADE'), primary_key=True)


class User(db.Model, UserMixin):
    """User model"""

    __bind_key__ = 'public'
    __table_args__ = {'schema': 'public'}

    # Versioning Information
    created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())

    # User Information
    id = db.Column(db.Integer, primary_key=True)
    user_name = db.Column(db.String(80), unique=True, nullable=False)
    email = db.Column(db.String(255), unique=True, nullable=False)
    password = db.Column(db.String(255), nullable=False)
    first_name = db.Column(db.String(80), nullable=False)
    last_name = db.Column(db.String(80), nullable=False)
    active = db.Column(db.Boolean)
    fs_uniquifier = db.Column(db.String(255), unique=True, nullable=False)
    confirmed_at = db.Column(db.DateTime, nullable=True)
    valid_to = db.Column(db.Date, nullable=True)

    # Security Trackable Information
    last_login_at = db.Column(db.DateTime, nullable=True)
    current_login_at = db.Column(db.DateTime, nullable=True)
    last_login_ip = db.Column(db.String(255), nullable=True)
    current_login_ip = db.Column(db.String(255), nullable=True)
    login_count = db.Column(db.Integer, nullable=False, default=0)

    # Relations
    roles = db.relationship('Role', secondary=RolesUsers.__table__, backref=db.backref('users', lazy='dynamic'))
    tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)

    def __repr__(self):
        return '<User %r>' % self.user_name

    def has_roles(self, *args):
        return any(role.name in args for role in self.roles)


class TenantDomain(db.Model):
    __bind_key__ = 'public'
    __table_args__ = {'schema': 'public'}

    id = db.Column(db.Integer, primary_key=True)
    tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
    # Originally, domain was required to be unique.
    # However, several tenants can run from the same domain (e.g. for demo purposes,
    # but also internal and external chat clients.
    domain = db.Column(db.String(255), nullable=False)
    valid_to = db.Column(db.Date, nullable=True)

    # Versioning Information
    created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
    created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=False)
    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
    updated_by = db.Column(db.Integer, db.ForeignKey(User.id))

    def __repr__(self):
        return f"<TenantDomain {self.id}: {self.domain}>"