- Improvement of Entitlements Domain
- Introduction of LicensePeriod - Introduction of Payments - Introduction of Invoices - Services definitions for Entitlements Domain
This commit is contained in:
Josako
@@ -2,14 +2,13 @@ import io
|
||||
import os
|
||||
from datetime import datetime as dt, timezone as tz, datetime
|
||||
|
||||
from celery import states
|
||||
from dateutil.relativedelta import relativedelta
|
||||
from flask import current_app
|
||||
from sqlalchemy import or_, and_, text
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from common.extensions import db
|
||||
from common.models.user import Tenant
|
||||
from common.models.entitlements import BusinessEventLog, LicenseUsage, License
|
||||
from common.services.entitlements.license_period_services import LicensePeriodServices
|
||||
from common.utils.celery_utils import current_celery
|
||||
from common.utils.eveai_exceptions import EveAINoLicenseForTenant, EveAIException, EveAINoActiveLicense
|
||||
from common.utils.database import Database
|
||||
@@ -21,52 +20,6 @@ def ping():
|
||||
return 'pong'
|
||||
|
||||
|
||||
@current_celery.task(name='update_usages', queue='entitlements')
|
||||
def update_usages():
|
||||
current_timestamp = dt.now(tz.utc)
|
||||
tenant_ids = get_all_tenant_ids()
|
||||
|
||||
# List to collect all errors
|
||||
error_list = []
|
||||
|
||||
for tenant_id in tenant_ids:
|
||||
if tenant_id == 1:
|
||||
continue
|
||||
try:
|
||||
Database(tenant_id).switch_schema()
|
||||
check_and_create_license_usage_for_tenant(tenant_id)
|
||||
tenant = Tenant.query.get(tenant_id)
|
||||
if tenant.storage_dirty:
|
||||
recalculate_storage_for_tenant(tenant)
|
||||
logs = get_logs_for_processing(tenant_id, current_timestamp)
|
||||
if not logs:
|
||||
continue # If no logs to be processed, continu to the next tenant
|
||||
|
||||
# Get the min and max timestamp from the logs
|
||||
min_timestamp = min(log.timestamp for log in logs)
|
||||
max_timestamp = max(log.timestamp for log in logs)
|
||||
|
||||
# Retrieve relevant LicenseUsage records
|
||||
license_usages = get_relevant_license_usages(db.session, tenant_id, min_timestamp, max_timestamp)
|
||||
|
||||
# Split logs based on LicenseUsage periods
|
||||
logs_by_usage = split_logs_by_license_usage(logs, license_usages)
|
||||
|
||||
# Now you can process logs for each LicenseUsage
|
||||
for license_usage_id, logs in logs_by_usage.items():
|
||||
process_logs_for_license_usage(tenant_id, license_usage_id, logs)
|
||||
except Exception as e:
|
||||
error = f"Usage Calculation error for Tenant {tenant_id}: {e}"
|
||||
error_list.append(error)
|
||||
current_app.logger.error(error)
|
||||
continue
|
||||
|
||||
if error_list:
|
||||
raise Exception('\n'.join(error_list))
|
||||
|
||||
return "Update Usages taks completed successfully"
|
||||
|
||||
|
||||
@current_celery.task(name='persist_business_events', queue='entitlements')
|
||||
def persist_business_events(log_entries):
|
||||
"""
|
||||
@@ -76,7 +29,7 @@ def persist_business_events(log_entries):
|
||||
log_entries: List of log event dictionaries to persist
|
||||
"""
|
||||
try:
|
||||
db_entries = []
|
||||
event_logs = []
|
||||
for entry in log_entries:
|
||||
event_log = BusinessEventLog(
|
||||
timestamp=entry.pop('timestamp'),
|
||||
@@ -103,119 +56,44 @@ def persist_business_events(log_entries):
|
||||
llm_interaction_type=entry.pop('llm_interaction_type', None),
|
||||
message=entry.pop('message', None)
|
||||
)
|
||||
db_entries.append(event_log)
|
||||
event_logs.append(event_log)
|
||||
|
||||
# Perform a bulk insert of all entries
|
||||
db.session.bulk_save_objects(db_entries)
|
||||
db.session.bulk_save_objects(event_logs)
|
||||
db.session.commit()
|
||||
current_app.logger.info(f"Successfully persisted {len(db_entries)} business event logs")
|
||||
current_app.logger.info(f"Successfully persisted {len(event_logs)} business event logs")
|
||||
|
||||
tenant_id = event_logs[0].tenant_id
|
||||
try:
|
||||
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
||||
except EveAIException as e:
|
||||
current_app.logger.error(f"Failed to find license period for tenant {tenant_id}: {str(e)}")
|
||||
return
|
||||
lic_usage = None
|
||||
if not license_period.license_usage:
|
||||
lic_usage = LicenseUsage(
|
||||
tenant_id=tenant_id,
|
||||
license_period_id=license_period.id,
|
||||
)
|
||||
try:
|
||||
db.session.add(lic_usage)
|
||||
db.session.commit()
|
||||
current_app.logger.info(f"Created new license usage for tenant {tenant_id}")
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error trying to create license usage for tenant {tenant_id}: {str(e)}")
|
||||
return
|
||||
else:
|
||||
lic_usage = license_period.license_usage
|
||||
|
||||
process_logs_for_license_usage(tenant_id, lic_usage, event_logs)
|
||||
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Failed to persist business event logs: {e}")
|
||||
db.session.rollback()
|
||||
|
||||
|
||||
def get_all_tenant_ids():
|
||||
tenant_ids = db.session.query(Tenant.id).all()
|
||||
return [tenant_id[0] for tenant_id in tenant_ids] # Extract tenant_id from tuples
|
||||
|
||||
|
||||
def check_and_create_license_usage_for_tenant(tenant_id):
|
||||
current_date = dt.now(tz.utc).date()
|
||||
license_usages = (db.session.query(LicenseUsage)
|
||||
.filter_by(tenant_id=tenant_id)
|
||||
.filter(and_(LicenseUsage.period_start_date <= current_date,
|
||||
LicenseUsage.period_end_date >= current_date))
|
||||
.all())
|
||||
if not license_usages:
|
||||
active_license = (db.session.query(License).filter_by(tenant_id=tenant_id)
|
||||
.filter(and_(License.start_date <= current_date,
|
||||
License.end_date >= current_date))
|
||||
.one_or_none())
|
||||
if not active_license:
|
||||
current_app.logger.error(f"No License defined for {tenant_id}. "
|
||||
f"Impossible to calculate license usage.")
|
||||
raise EveAINoActiveLicense(tenant_id)
|
||||
|
||||
start_date, end_date = calculate_valid_period(current_date, active_license.start_date)
|
||||
new_license_usage = LicenseUsage(period_start_date=start_date,
|
||||
period_end_date=end_date,
|
||||
license_id=active_license.id,
|
||||
tenant_id=tenant_id
|
||||
)
|
||||
try:
|
||||
db.session.add(new_license_usage)
|
||||
db.session.commit()
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error trying to create new license usage for tenant {tenant_id}. "
|
||||
f"Error: {str(e)}")
|
||||
raise e
|
||||
|
||||
|
||||
def calculate_valid_period(given_date, original_start_date):
|
||||
# Ensure both dates are of datetime.date type
|
||||
if isinstance(given_date, datetime):
|
||||
given_date = given_date.date()
|
||||
if isinstance(original_start_date, datetime):
|
||||
original_start_date = original_start_date.date()
|
||||
|
||||
# Step 1: Find the most recent start_date less than or equal to given_date
|
||||
start_date = original_start_date
|
||||
while start_date <= given_date:
|
||||
next_start_date = start_date + relativedelta(months=1)
|
||||
if next_start_date > given_date:
|
||||
break
|
||||
start_date = next_start_date
|
||||
|
||||
# Step 2: Calculate the end_date for this period
|
||||
end_date = start_date + relativedelta(months=1, days=-1)
|
||||
|
||||
# Ensure the given date falls within the period
|
||||
if start_date <= given_date <= end_date:
|
||||
return start_date, end_date
|
||||
else:
|
||||
raise ValueError("Given date does not fall within a valid period.")
|
||||
|
||||
|
||||
def get_logs_for_processing(tenant_id, end_time_stamp):
|
||||
return (db.session.query(BusinessEventLog).filter(
|
||||
BusinessEventLog.tenant_id == tenant_id,
|
||||
BusinessEventLog.license_usage_id == None,
|
||||
BusinessEventLog.timestamp <= end_time_stamp,
|
||||
).all())
|
||||
|
||||
|
||||
def get_relevant_license_usages(session, tenant_id, min_timestamp, max_timestamp):
|
||||
# Fetch LicenseUsage records where the log timestamps fall between period_start_date and period_end_date
|
||||
return session.query(LicenseUsage).filter(
|
||||
LicenseUsage.tenant_id == tenant_id,
|
||||
LicenseUsage.period_start_date <= max_timestamp.date(),
|
||||
LicenseUsage.period_end_date >= min_timestamp.date()
|
||||
).order_by(LicenseUsage.period_start_date).all()
|
||||
|
||||
|
||||
def split_logs_by_license_usage(logs, license_usages):
|
||||
# Dictionary to hold logs categorized by LicenseUsage
|
||||
logs_by_usage = {lu.id: [] for lu in license_usages}
|
||||
|
||||
for log in logs:
|
||||
# Find the corresponding LicenseUsage for each log based on the timestamp
|
||||
for license_usage in license_usages:
|
||||
if license_usage.period_start_date <= log.timestamp.date() <= license_usage.period_end_date:
|
||||
logs_by_usage[license_usage.id].append(log)
|
||||
break
|
||||
|
||||
return logs_by_usage
|
||||
|
||||
|
||||
def process_logs_for_license_usage(tenant_id, license_usage_id, logs):
|
||||
# Retrieve the LicenseUsage record
|
||||
license_usage = db.session.query(LicenseUsage).filter_by(id=license_usage_id).first()
|
||||
|
||||
if not license_usage:
|
||||
raise ValueError(f"LicenseUsage with id {license_usage_id} not found.")
|
||||
|
||||
def process_logs_for_license_usage(tenant_id, license_usage, logs):
|
||||
# Initialize variables to accumulate usage data
|
||||
embedding_mb_used = 0
|
||||
embedding_prompt_tokens_used = 0
|
||||
@@ -225,10 +103,13 @@ def process_logs_for_license_usage(tenant_id, license_usage_id, logs):
|
||||
interaction_completion_tokens_used = 0
|
||||
interaction_total_tokens_used = 0
|
||||
|
||||
recalculate_storage = False
|
||||
|
||||
# Process each log
|
||||
for log in logs:
|
||||
# Case for 'Create Embeddings' event
|
||||
if log.event_type == 'Create Embeddings':
|
||||
recalculate_storage = True
|
||||
if log.message == 'Starting Trace for Create Embeddings':
|
||||
embedding_mb_used += log.document_version_file_size
|
||||
elif log.message == 'Final LLM Metrics':
|
||||
@@ -256,7 +137,7 @@ def process_logs_for_license_usage(tenant_id, license_usage_id, logs):
|
||||
interaction_total_tokens_used += log.llm_metrics_total_tokens
|
||||
|
||||
# Mark the log as processed by setting the license_usage_id
|
||||
log.license_usage_id = license_usage_id
|
||||
log.license_usage_id = license_usage.id
|
||||
|
||||
# Update the LicenseUsage record with the accumulated values
|
||||
license_usage.embedding_mb_used += embedding_mb_used
|
||||
@@ -267,6 +148,9 @@ def process_logs_for_license_usage(tenant_id, license_usage_id, logs):
|
||||
license_usage.interaction_completion_tokens_used += interaction_completion_tokens_used
|
||||
license_usage.interaction_total_tokens_used += interaction_total_tokens_used
|
||||
|
||||
if recalculate_storage:
|
||||
recalculate_storage_for_tenant(tenant_id)
|
||||
|
||||
# Commit the updates to the LicenseUsage and log records
|
||||
try:
|
||||
db.session.add(license_usage)
|
||||
@@ -279,7 +163,8 @@ def process_logs_for_license_usage(tenant_id, license_usage_id, logs):
|
||||
raise e
|
||||
|
||||
|
||||
def recalculate_storage_for_tenant(tenant):
|
||||
def recalculate_storage_for_tenant(tenant_id):
|
||||
Database(tenant_id).switch_schema()
|
||||
# Perform a SUM operation to get the total file size from document_versions
|
||||
total_storage = db.session.execute(text(f"""
|
||||
SELECT SUM(file_size)
|
||||
@@ -287,19 +172,15 @@ def recalculate_storage_for_tenant(tenant):
|
||||
""")).scalar()
|
||||
|
||||
# Update the LicenseUsage with the recalculated storage
|
||||
license_usage = db.session.query(LicenseUsage).filter_by(tenant_id=tenant.id).first()
|
||||
license_usage = db.session.query(LicenseUsage).filter_by(tenant_id=tenant_id).first()
|
||||
license_usage.storage_mb_used = total_storage
|
||||
|
||||
# Reset the dirty flag after recalculating
|
||||
tenant.storage_dirty = False
|
||||
|
||||
# Commit the changes
|
||||
try:
|
||||
db.session.add(tenant)
|
||||
db.session.add(license_usage)
|
||||
db.session.commit()
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error trying to update tenant {tenant.id} for Dirty Storage. ")
|
||||
current_app.logger.error(f"Error trying to update tenant {tenant_id} for Dirty Storage. ")
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user