144 lines
6.2 KiB
Python
144 lines
6.2 KiB
Python
from dateutil.relativedelta import relativedelta
|
|
from flask import session, current_app, flash
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from sqlalchemy.sql.expression import text
|
|
|
|
from common.extensions import db, cache_manager
|
|
from common.models.entitlements import PartnerServiceLicenseTier, License, LicenseUsage, LicensePeriod, PeriodStatus
|
|
from common.models.user import Partner, PartnerTenant
|
|
from common.services.entitlements import LicensePeriodServices
|
|
from common.utils.database import Database
|
|
from common.utils.eveai_exceptions import EveAINoManagementPartnerService, EveAINoActiveLicense, \
|
|
EveAIStorageQuotaExceeded, EveAIEmbeddingQuotaExceeded, EveAIInteractionQuotaExceeded, EveAILicensePeriodsExceeded, \
|
|
EveAIException
|
|
from common.utils.model_logging_utils import set_logging_information, update_logging_information
|
|
from datetime import datetime as dt, timezone as tz
|
|
|
|
from common.utils.security_utils import current_user_has_role
|
|
|
|
|
|
class LicenseUsageServices:
|
|
@staticmethod
|
|
def check_storage_and_embedding_quota(tenant_id: int, file_size_mb: float) -> None:
|
|
"""
|
|
Check if a tenant can add a new document without exceeding storage and embedding quotas
|
|
|
|
Args:
|
|
tenant_id: ID of the tenant
|
|
file_size_mb: Size of the file in MB
|
|
|
|
Raises:
|
|
EveAIStorageQuotaExceeded: If storage quota would be exceeded
|
|
EveAIEmbeddingQuotaExceeded: If embedding quota would be exceeded
|
|
EveAINoActiveLicense: If no active license is found
|
|
EveAIException: For other errors
|
|
"""
|
|
# Get active license period
|
|
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
|
# Early return if both overruns are allowed - no need to check usage at all
|
|
if license_period.additional_storage_allowed and license_period.additional_embedding_allowed:
|
|
return
|
|
|
|
# Check storage quota only if overruns are not allowed
|
|
if not license_period.additional_storage_allowed:
|
|
LicenseUsageServices._validate_storage_quota(license_period, file_size_mb)
|
|
|
|
# Check embedding quota only if overruns are not allowed
|
|
if not license_period.additional_embedding_allowed:
|
|
LicenseUsageServices._validate_embedding_quota(license_period, file_size_mb)
|
|
|
|
@staticmethod
|
|
def check_embedding_quota(tenant_id: int, file_size_mb: float) -> None:
|
|
"""
|
|
Check if a tenant can re-embed a document without exceeding embedding quota
|
|
|
|
Args:
|
|
tenant_id: ID of the tenant
|
|
file_size_mb: Size of the file in MB
|
|
|
|
Raises:
|
|
EveAIEmbeddingQuotaExceeded: If embedding quota would be exceeded
|
|
EveAINoActiveLicense: If no active license is found
|
|
EveAIException: For other errors
|
|
"""
|
|
# Get active license period
|
|
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
|
# Early return if both overruns are allowed - no need to check usage at all
|
|
if license_period.additional_embedding_allowed:
|
|
return
|
|
|
|
# Check embedding quota
|
|
LicenseUsageServices._validate_embedding_quota(license_period, file_size_mb)
|
|
|
|
@staticmethod
|
|
def check_interaction_quota(tenant_id: int) -> None:
|
|
"""
|
|
Check if a tenant can execute a specialist without exceeding interaction quota. As it is impossible to estimate
|
|
the number of interaction tokens, we only check if the interaction quota are exceeded. So we might have a
|
|
limited overrun.
|
|
|
|
Args:
|
|
tenant_id: ID of the tenant
|
|
|
|
Raises:
|
|
EveAIInteractionQuotaExceeded: If interaction quota would be exceeded
|
|
EveAINoActiveLicense: If no active license is found
|
|
EveAIException: For other errors
|
|
"""
|
|
# Get active license period
|
|
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
|
# Early return if both overruns are allowed - no need to check usage at all
|
|
if license_period.additional_interaction_allowed:
|
|
return
|
|
|
|
# Convert tokens to M tokens and check interaction quota
|
|
LicenseUsageServices._validate_interaction_quota(license_period)
|
|
|
|
@staticmethod
|
|
def _validate_storage_quota(license_period: LicensePeriod, additional_mb: float) -> None:
|
|
"""Check storage quota and raise exception if exceeded"""
|
|
current_storage = license_period.license_usage.storage_mb_used or 0
|
|
projected_storage = current_storage + additional_mb
|
|
max_storage = license_period.max_storage_mb
|
|
|
|
# Hard limit check (we only get here if overruns are NOT allowed)
|
|
if projected_storage > max_storage:
|
|
raise EveAIStorageQuotaExceeded(
|
|
current_usage=current_storage,
|
|
limit=max_storage,
|
|
additional=additional_mb
|
|
)
|
|
|
|
@staticmethod
|
|
def _validate_embedding_quota(license_period: LicensePeriod, additional_mb: float) -> None:
|
|
"""Check embedding quota and raise exception if exceeded"""
|
|
current_embedding = license_period.license_usage.embedding_mb_used or 0
|
|
projected_embedding = current_embedding + additional_mb
|
|
max_embedding = license_period.included_embedding_mb
|
|
|
|
# Hard limit check (we only get here if overruns are NOT allowed)
|
|
if projected_embedding > max_embedding:
|
|
raise EveAIEmbeddingQuotaExceeded(
|
|
current_usage=current_embedding,
|
|
limit=max_embedding,
|
|
additional=additional_mb
|
|
)
|
|
|
|
@staticmethod
|
|
def _validate_interaction_quota(license_period) -> None:
|
|
"""Check interaction quota and raise exception if exceeded (tokens in millions). We might have an overrun!"""
|
|
current_tokens = license_period.license_usage.interaction_total_tokens_used / 1_000_000 or 0
|
|
max_tokens = license_period.included_interaction_tokens
|
|
|
|
# Hard limit check (we only get here if overruns are NOT allowed)
|
|
if current_tokens > max_tokens:
|
|
raise EveAIInteractionQuotaExceeded(
|
|
current_usage=current_tokens,
|
|
limit=max_tokens
|
|
)
|
|
|
|
|
|
|
|
|
|
|