- Minor corrections to entitlement changes and upgrades

- started new eveai_entitlements component (not finished)

common/models/entitlements.py

@@ -95,10 +95,13 @@ class LicenseUsage(db.Model):
     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)
+    embedding_prompt_tokens_used = db.Column(db.Integer, default=0)
+    embedding_completion_tokens_used = db.Column(db.Integer, default=0)
+    embedding_total_tokens_used = db.Column(db.Integer, default=0)
+    interaction_prompt_tokens_used = db.Column(db.Integer, default=0)
+    interaction_completion_tokens_used = db.Column(db.Integer, default=0)
+    interaction_total_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)
 

common/models/user.py

@@ -64,6 +64,7 @@ class Tenant(db.Model):
     # 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')

common/utils/document_utils.py

@@ -12,6 +12,7 @@ import requests
 from urllib.parse import urlparse, unquote
 import os
 from .eveai_exceptions import EveAIInvalidLanguageException, EveAIDoubleURLException, EveAIUnsupportedFileType
+from ..models.user import Tenant
 
 
 def create_document_stack(api_input, file, filename, extension, tenant_id):
@@ -81,6 +82,8 @@ def create_version_for_document(document, url, language, user_context, user_meta
 
     set_logging_information(new_doc_vers, dt.now(tz.utc))
 
+    mark_tenant_storage_dirty(document.tenant_id)
+
     return new_doc_vers
 
 
@@ -338,3 +341,10 @@ def refresh_document(doc_id):
     }
 
     return refresh_document_with_info(doc_id, api_input)
+
+
+# Function triggered when a document_version is created or updated
+def mark_tenant_storage_dirty(tenant_id):
+    tenant = db.session.query(Tenant).filter_by(id=tenant_id).first()
+    tenant.storage_dirty = True
+    db.session.commit()

common/utils/eveai_exceptions.py

@@ -34,3 +34,10 @@ class EveAIUnsupportedFileType(EveAIException):
         super().__init__(message, status_code, payload)
 
 
+class EveAINoLicenseForTenant(EveAIException):
+    """Raised when no active license for a tenant is provided"""
+
+    def __init__(self, message="No license for tenant found", status_code=400, payload=None):
+        super().__init__(message, status_code, payload)
+
+