- 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

common/utils/business_event.py

@@ -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,

common/utils/document_utils.py

@@ -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.')

common/utils/minio_utils.py

@@ -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}")
 

common/utils/view_assistants.py

@@ -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):