Files
eveAI/common/models/user.py
Josako d6a7743f26 - Minor corrections to entitlement changes and upgrades
- started new eveai_entitlements component (not finished)
2024-10-08 09:12:16 +02:00

196 lines
7.9 KiB
Python

from datetime import date
from common.extensions import db
from flask_security import UserMixin, RoleMixin
from sqlalchemy.dialects.postgresql import ARRAY
import sqlalchemy as sa
from common.models.entitlements import License
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)
timezone = db.Column(db.String(50), nullable=True, default='UTC')
rag_context = db.Column(db.Text, nullable=True)
type = db.Column(db.String(20), nullable=True, server_default='Active')
# 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)
html_excluded_classes = db.Column(ARRAY(sa.String(200)), nullable=True)
min_chunk_size = db.Column(db.Integer, nullable=True, default=2000)
max_chunk_size = db.Column(db.Integer, nullable=True, default=3000)
# 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)
fallback_algorithms = db.Column(ARRAY(sa.String(50)), nullable=True)
# Licensing Information
encrypted_chat_api_key = db.Column(db.String(500), nullable=True)
encrypted_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)
# Entitlements
currency = db.Column(db.String(20), nullable=True)
usage_email = db.Column(db.String(255), nullable=True)
storage_dirty = db.Column(db.Boolean, nullable=True, default=False)
# Relations
users = db.relationship('User', backref='tenant')
domains = db.relationship('TenantDomain', backref='tenant')
licenses = db.relationship('License', back_populates='tenant')
license_usages = db.relationship('LicenseUsage', backref='tenant')
@property
def current_license(self):
today = date.today()
return License.query.filter(
License.tenant_id == self.id,
License.start_date <= today,
(License.end_date.is_(None) | (License.end_date >= today))
).order_by(License.start_date.desc()).first()
def __repr__(self):
return f"<Tenant {self.id}: {self.name}>"
def to_dict(self):
return {
'id': self.id,
'name': self.name,
'website': self.website,
'timezone': self.timezone,
'rag_context': self.rag_context,
'type': self.type,
'default_language': self.default_language,
'allowed_languages': self.allowed_languages,
'embedding_model': self.embedding_model,
'llm_model': self.llm_model,
'html_tags': self.html_tags,
'html_end_tags': self.html_end_tags,
'html_included_elements': self.html_included_elements,
'html_excluded_elements': self.html_excluded_elements,
'html_excluded_classes': self.html_excluded_classes,
'min_chunk_size': self.min_chunk_size,
'max_chunk_size': self.max_chunk_size,
'es_k': self.es_k,
'es_similarity_threshold': self.es_similarity_threshold,
'chat_RAG_temperature': self.chat_RAG_temperature,
'chat_no_RAG_temperature': self.chat_no_RAG_temperature,
'fallback_algorithms': self.fallback_algorithms,
'embed_tuning': self.embed_tuning,
'rag_tuning': self.rag_tuning,
'currency': self.currency,
'usage_email': self.usage_email,
}
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=True)
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}>"