- Refined entitlements to work with MiB for both embeddings and storage

- Improved DocumentVersion storage attributes to reflect Minio settings
- Added size to DocumentVersions to easily calculate usage
- License / LicenseTier forms and views added
This commit is contained in:
Josako
2024-10-07 14:17:44 +02:00
parent f638860e90
commit 9782e31ae5
31 changed files with 1416 additions and 83 deletions

View File

@@ -28,9 +28,10 @@ class DocumentVersion(db.Model):
id = db.Column(db.Integer, primary_key=True)
doc_id = db.Column(db.Integer, db.ForeignKey(Document.id), nullable=False)
url = db.Column(db.String(200), nullable=True)
file_location = db.Column(db.String(255), nullable=True)
file_name = db.Column(db.String(200), nullable=True)
bucket_name = db.Column(db.String(255), nullable=True)
object_name = db.Column(db.String(200), nullable=True)
file_type = db.Column(db.String(20), nullable=True)
file_size = db.Column(db.Float, nullable=True)
language = db.Column(db.String(2), nullable=False)
user_context = db.Column(db.Text, nullable=True)
system_context = db.Column(db.Text, nullable=True)

View File

@@ -0,0 +1,107 @@
from common.extensions import db
class BusinessEventLog(db.Model):
__bind_key__ = 'public'
__table_args__ = {'schema': 'public'}
id = db.Column(db.Integer, primary_key=True)
timestamp = db.Column(db.DateTime, nullable=False)
event_type = db.Column(db.String(50), nullable=False)
tenant_id = db.Column(db.Integer, nullable=False)
trace_id = db.Column(db.String(50), nullable=False)
span_id = db.Column(db.String(50))
span_name = db.Column(db.String(50))
parent_span_id = db.Column(db.String(50))
document_version_id = db.Column(db.Integer)
document_version_file_size = db.Column(db.Float)
chat_session_id = db.Column(db.String(50))
interaction_id = db.Column(db.Integer)
environment = db.Column(db.String(20))
llm_metrics_total_tokens = db.Column(db.Integer)
llm_metrics_prompt_tokens = db.Column(db.Integer)
llm_metrics_completion_tokens = db.Column(db.Integer)
llm_metrics_total_time = db.Column(db.Float)
llm_metrics_call_count = db.Column(db.Integer)
llm_interaction_type = db.Column(db.String(20))
message = db.Column(db.Text)
license_usage_id = db.Column(db.Integer, db.ForeignKey('public.license_usage.id'), nullable=True)
license_usage = db.relationship('LicenseUsage', backref='events')
class License(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)
tier_id = db.Column(db.Integer, db.ForeignKey('public.license_tier.id'),nullable=False) # 'small', 'medium', 'custom'
start_date = db.Column(db.Date, nullable=False)
end_date = db.Column(db.Date, nullable=True)
currency = db.Column(db.String(20), nullable=False)
yearly_payment = db.Column(db.Boolean, nullable=False, default=False)
basic_fee = db.Column(db.Float, nullable=False)
max_storage_mb = db.Column(db.Integer, nullable=False)
additional_storage_price = db.Column(db.Float, nullable=False)
additional_storage_bucket = db.Column(db.Integer, nullable=False)
included_embedding_mb = db.Column(db.Integer, nullable=False)
additional_embedding_price = db.Column(db.Numeric(10, 4), nullable=False)
additional_embedding_bucket = db.Column(db.Integer, nullable=False)
included_interaction_tokens = db.Column(db.Integer, nullable=False)
additional_interaction_token_price = db.Column(db.Numeric(10, 4), nullable=False)
additional_interaction_bucket = db.Column(db.Integer, nullable=False)
overage_embedding = db.Column(db.Float, nullable=False, default=0)
overage_interaction = db.Column(db.Float, nullable=False, default=0)
tenant = db.relationship('Tenant', back_populates='licenses')
license_tier = db.relationship('LicenseTier', back_populates='licenses')
usages = db.relationship('LicenseUsage', order_by='LicenseUsage.period_start_date', back_populates='license')
class LicenseTier(db.Model):
__bind_key__ = 'public'
__table_args__ = {'schema': 'public'}
id = db.Column(db.Integer, primary_key=True)
name = db.Column(db.String(50), nullable=False)
version = db.Column(db.String(50), nullable=False)
start_date = db.Column(db.Date, nullable=False)
end_date = db.Column(db.Date, nullable=True)
basic_fee_d = db.Column(db.Float, nullable=True)
basic_fee_e = db.Column(db.Float, nullable=True)
max_storage_mb = db.Column(db.Integer, nullable=False)
additional_storage_price_d = db.Column(db.Numeric(10, 4), nullable=False)
additional_storage_price_e = db.Column(db.Numeric(10, 4), nullable=False)
additional_storage_bucket = db.Column(db.Integer, nullable=False)
included_embedding_mb = db.Column(db.Integer, nullable=False)
additional_embedding_price_d = db.Column(db.Numeric(10, 4), nullable=False)
additional_embedding_price_e = db.Column(db.Numeric(10, 4), nullable=False)
additional_embedding_bucket = db.Column(db.Integer, nullable=False)
included_interaction_tokens = db.Column(db.Integer, nullable=False)
additional_interaction_token_price_d = db.Column(db.Numeric(10, 4), nullable=False)
additional_interaction_token_price_e = db.Column(db.Numeric(10, 4), nullable=False)
additional_interaction_bucket = db.Column(db.Integer, nullable=False)
standard_overage_embedding = db.Column(db.Float, nullable=False, default=0)
standard_overage_interaction = db.Column(db.Float, nullable=False, default=0)
licenses = db.relationship('License', back_populates='license_tier')
class LicenseUsage(db.Model):
__bind_key__ = 'public'
__table_args__ = {'schema': 'public'}
id = db.Column(db.Integer, primary_key=True)
license_id = db.Column(db.Integer, db.ForeignKey('public.license.id'), nullable=False)
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
storage_mb_used = db.Column(db.Integer, default=0)
storage_tokens_used = db.Column(db.Integer, default=0)
embedding_mb_used = db.Column(db.Integer, default=0)
embedding_tokens_used = db.Column(db.Integer, default=0)
interaction_tokens_used = db.Column(db.Integer, default=0)
period_start_date = db.Column(db.Date, nullable=False)
period_end_date = db.Column(db.Date, nullable=False)
license = db.relationship('License', back_populates='usages')

View File

@@ -1,27 +0,0 @@
from common.extensions import db
class BusinessEventLog(db.Model):
__bind_key__ = 'public'
__table_args__ = {'schema': 'public'}
id = db.Column(db.Integer, primary_key=True)
timestamp = db.Column(db.DateTime, nullable=False)
event_type = db.Column(db.String(50), nullable=False)
tenant_id = db.Column(db.Integer, nullable=False)
trace_id = db.Column(db.String(50), nullable=False)
span_id = db.Column(db.String(50))
span_name = db.Column(db.String(50))
parent_span_id = db.Column(db.String(50))
document_version_id = db.Column(db.Integer)
chat_session_id = db.Column(db.String(50))
interaction_id = db.Column(db.Integer)
environment = db.Column(db.String(20))
llm_metrics_total_tokens = db.Column(db.Integer)
llm_metrics_prompt_tokens = db.Column(db.Integer)
llm_metrics_completion_tokens = db.Column(db.Integer)
llm_metrics_total_time = db.Column(db.Float)
llm_metrics_call_count = db.Column(db.Integer)
llm_interaction_type = db.Column(db.String(20))
message = db.Column(db.Text)
# Add any other fields relevant for invoicing or warnings

View File

@@ -1,8 +1,12 @@
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"""
@@ -50,9 +54,6 @@ class Tenant(db.Model):
fallback_algorithms = db.Column(ARRAY(sa.String(50)), nullable=True)
# 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)
encrypted_api_key = db.Column(db.String(500), nullable=True)
@@ -60,9 +61,24 @@ class Tenant(db.Model):
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)
# 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}>"
@@ -91,11 +107,10 @@ class Tenant(db.Model):
'chat_RAG_temperature': self.chat_RAG_temperature,
'chat_no_RAG_temperature': self.chat_no_RAG_temperature,
'fallback_algorithms': self.fallback_algorithms,
'license_start_date': self.license_start_date,
'license_end_date': self.license_end_date,
'allowed_monthly_interactions': self.allowed_monthly_interactions,
'embed_tuning': self.embed_tuning,
'rag_tuning': self.rag_tuning,
'currency': self.currency,
'usage_email': self.usage_email,
}

View File

@@ -8,7 +8,7 @@ from portkey_ai import Portkey, Config
import logging
from .business_event_context import BusinessEventContext
from common.models.monitoring import BusinessEventLog
from common.models.entitlements import BusinessEventLog
from common.extensions import db
@@ -25,6 +25,7 @@ class BusinessEvent:
self.span_name = None
self.parent_span_id = None
self.document_version_id = kwargs.get('document_version_id')
self.document_version_file_size = kwargs.get('document_version_file_size')
self.chat_session_id = kwargs.get('chat_session_id')
self.interaction_id = kwargs.get('interaction_id')
self.environment = os.environ.get("FLASK_ENV", "development")
@@ -107,6 +108,7 @@ class BusinessEvent:
'span_name': self.span_name,
'parent_span_id': self.parent_span_id,
'document_version_id': self.document_version_id,
'document_version_file_size': self.document_version_file_size,
'chat_session_id': self.chat_session_id,
'interaction_id': self.interaction_id,
'environment': self.environment,
@@ -124,6 +126,7 @@ class BusinessEvent:
span_name=self.span_name,
parent_span_id=self.parent_span_id,
document_version_id=self.document_version_id,
document_version_file_size=self.document_version_file_size,
chat_session_id=self.chat_session_id,
interaction_id=self.interaction_id,
environment=self.environment,
@@ -144,6 +147,7 @@ class BusinessEvent:
'span_name': self.span_name,
'parent_span_id': self.parent_span_id,
'document_version_id': self.document_version_id,
'document_version_file_size': self.document_version_file_size,
'chat_session_id': self.chat_session_id,
'interaction_id': self.interaction_id,
'environment': self.environment,
@@ -166,6 +170,7 @@ class BusinessEvent:
span_name=self.span_name,
parent_span_id=self.parent_span_id,
document_version_id=self.document_version_id,
document_version_file_size=self.document_version_file_size,
chat_session_id=self.chat_session_id,
interaction_id=self.interaction_id,
environment=self.environment,
@@ -190,6 +195,7 @@ class BusinessEvent:
'span_name': self.span_name,
'parent_span_id': self.parent_span_id,
'document_version_id': self.document_version_id,
'document_version_file_size': self.document_version_file_size,
'chat_session_id': self.chat_session_id,
'interaction_id': self.interaction_id,
'environment': self.environment,
@@ -213,6 +219,7 @@ class BusinessEvent:
span_name=self.span_name,
parent_span_id=self.parent_span_id,
document_version_id=self.document_version_id,
document_version_file_size=self.document_version_file_size,
chat_session_id=self.chat_session_id,
interaction_id=self.interaction_id,
environment=self.environment,

View File

@@ -86,14 +86,12 @@ def create_version_for_document(document, url, language, user_context, user_meta
def upload_file_for_version(doc_vers, file, extension, tenant_id):
doc_vers.file_type = extension
doc_vers.file_name = doc_vers.calc_file_name()
doc_vers.file_location = doc_vers.calc_file_location()
# Normally, the tenant bucket should exist. But let's be on the safe side if a migration took place.
minio_client.create_tenant_bucket(tenant_id)
try:
minio_client.upload_document_file(
bn, on, size = minio_client.upload_document_file(
tenant_id,
doc_vers.doc_id,
doc_vers.language,
@@ -101,6 +99,10 @@ def upload_file_for_version(doc_vers, file, extension, tenant_id):
doc_vers.file_name,
file
)
doc_vers.bucket_name = bn
doc_vers.object_name = on
doc_vers.file_size_mb = size / 1048576 # Convert bytes to MB
db.session.commit()
current_app.logger.info(f'Successfully saved document to MinIO for tenant {tenant_id} for '
f'document version {doc_vers.id} while uploading file.')

View File

@@ -50,7 +50,7 @@ class MinioClient:
self.client.put_object(
bucket_name, object_name, io.BytesIO(file_data), len(file_data)
)
return True
return bucket_name, object_name, len(file_data)
except S3Error as err:
raise Exception(f"Error occurred while uploading file: {err}")

View File

@@ -44,7 +44,7 @@ def form_validation_failed(request, form):
for fieldName, errorMessages in form.errors.items():
for err in errorMessages:
flash(f"Error in {fieldName}: {err}", 'danger')
current_app.logger.debug(f"Error in {fieldName}: {err}", 'danger')
current_app.logger.debug(f"Error in {fieldName}: {err}")
def form_to_dict(form):