pieter/eveAI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

40
migrations/public/versions/48714f1baac5_add_entitlement_information_to_tenant.py Normal file
View File

@@ -0,0 +1,40 @@
"""Add entitlement information to Tenant
Revision ID: 48714f1baac5
Revises: f201bfd23152
Create Date: 2024-10-03 14:49:53.922320
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = '48714f1baac5'
down_revision = 'f201bfd23152'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.add_column(sa.Column('yearly_payment', sa.Boolean(), nullable=False))
with op.batch_alter_table('tenant', schema=None) as batch_op:
batch_op.add_column(sa.Column('currency', sa.String(length=20), nullable=True))
batch_op.add_column(sa.Column('usage_email', sa.String(length=255), nullable=True))
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('tenant', schema=None) as batch_op:
batch_op.drop_column('usage_email')
batch_op.drop_column('currency')
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.drop_column('yearly_payment')
# ### end Alembic commands ###

106
migrations/public/versions/560d08d91e5b_introducing_new_licensing_approach.py Normal file
View File

@@ -0,0 +1,106 @@
"""Introducing new licensing approach
Revision ID: 560d08d91e5b
Revises: 254932fe7fe3
Create Date: 2024-10-02 15:29:05.963865
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = '560d08d91e5b'
down_revision = '254932fe7fe3'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('license_tier',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('name', sa.String(length=50), nullable=False),
sa.Column('version', sa.String(length=50), nullable=False),
sa.Column('start_date', sa.Date(), nullable=False),
sa.Column('end_date', sa.Date(), nullable=True),
sa.Column('basic_fee_d', sa.Float(), nullable=True),
sa.Column('basic_fee_e', sa.Float(), nullable=True),
sa.Column('max_storage_tokens', sa.Integer(), nullable=False),
sa.Column('additional_storage_token_price_d', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_storage_token_price_e', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('included_embedding_tokens', sa.Integer(), nullable=False),
sa.Column('additional_embedding_token_price_d', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_embedding_token_price_e', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_embedding_bucket', sa.Integer(), nullable=False),
sa.Column('included_interaction_tokens', sa.Integer(), nullable=False),
sa.Column('additional_interaction_token_price_d', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_interaction_token_price_e', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_interaction_bucket', sa.Integer(), nullable=False),
sa.Column('allow_overage', sa.Boolean(), nullable=True),
sa.PrimaryKeyConstraint('id'),
schema='public'
)
op.create_table('license',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('tenant_id', sa.Integer(), nullable=False),
sa.Column('tier_id', sa.Integer(), nullable=False),
sa.Column('start_date', sa.Date(), nullable=False),
sa.Column('end_date', sa.Date(), nullable=True),
sa.Column('currency', sa.String(length=20), nullable=False),
sa.Column('basic_fee', sa.Float(), nullable=False),
sa.Column('max_storage_tokens', sa.Integer(), nullable=False),
sa.Column('additional_storage_token_price', sa.Float(), nullable=False),
sa.Column('included_embedding_tokens', sa.Integer(), nullable=False),
sa.Column('additional_embedding_token_price', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_embedding_bucket', sa.Integer(), nullable=False),
sa.Column('included_interaction_tokens', sa.Integer(), nullable=False),
sa.Column('additional_interaction_token_price', sa.Numeric(precision=10, scale=4), nullable=False),
sa.Column('additional_interaction_bucket', sa.Integer(), nullable=False),
sa.Column('allow_overage', sa.Boolean(), nullable=True),
sa.ForeignKeyConstraint(['tenant_id'], ['public.tenant.id'], ),
sa.ForeignKeyConstraint(['tier_id'], ['public.license_tier.id'], ),
sa.PrimaryKeyConstraint('id'),
schema='public'
)
op.create_table('license_usage',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('license_id', sa.Integer(), nullable=False),
sa.Column('tenant_id', sa.Integer(), nullable=False),
sa.Column('storage_tokens_used', sa.Integer(), nullable=True),
sa.Column('embedding_tokens_used', sa.Integer(), nullable=True),
sa.Column('interaction_tokens_used', sa.Integer(), nullable=True),
sa.Column('period_start_date', sa.Date(), nullable=False),
sa.Column('period_end_date', sa.Date(), nullable=False),
sa.ForeignKeyConstraint(['license_id'], ['public.license.id'], ),
sa.ForeignKeyConstraint(['tenant_id'], ['public.tenant.id'], ),
sa.PrimaryKeyConstraint('id'),
schema='public'
)
with op.batch_alter_table('business_event_log', schema=None) as batch_op:
batch_op.add_column(sa.Column('license_usage_id', sa.Integer(), nullable=True))
batch_op.create_foreign_key(None, 'license_usage', ['license_usage_id'], ['id'], referent_schema='public')
with op.batch_alter_table('tenant', schema=None) as batch_op:
batch_op.drop_column('license_end_date')
batch_op.drop_column('allowed_monthly_interactions')
batch_op.drop_column('license_start_date')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('tenant', schema=None) as batch_op:
batch_op.add_column(sa.Column('license_start_date', sa.DATE(), autoincrement=False, nullable=True))
batch_op.add_column(sa.Column('allowed_monthly_interactions', sa.INTEGER(), autoincrement=False, nullable=True))
batch_op.add_column(sa.Column('license_end_date', sa.DATE(), autoincrement=False, nullable=True))
with op.batch_alter_table('business_event_log', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.drop_column('license_usage_id')
op.drop_table('license_usage', schema='public')
op.drop_table('license', schema='public')
op.drop_table('license_tier', schema='public')
# ### end Alembic commands ###

38
migrations/public/versions/6a7743d08106_adapt_licenseusage_to_accept_mb_iso_.py Normal file
View File

@@ -0,0 +1,38 @@
"""Adapt LicenseUsage to accept mb iso tokoens for Storage and Embeddings
Revision ID: 6a7743d08106
Revises: 9429f244f1a5
Create Date: 2024-10-07 06:04:39.424243
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = '6a7743d08106'
down_revision = '9429f244f1a5'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license_usage', schema=None) as batch_op:
batch_op.add_column(sa.Column('storage_mb_used', sa.Integer(), nullable=True))
batch_op.add_column(sa.Column('embedding_mb_used', sa.Integer(), nullable=True))
batch_op.drop_column('storage_tokens_used')
batch_op.drop_column('embedding_tokens_used')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license_usage', schema=None) as batch_op:
batch_op.add_column(sa.Column('embedding_tokens_used', sa.INTEGER(), autoincrement=False, nullable=True))
batch_op.add_column(sa.Column('storage_tokens_used', sa.INTEGER(), autoincrement=False, nullable=True))
batch_op.drop_column('embedding_mb_used')
batch_op.drop_column('storage_mb_used')
# ### end Alembic commands ###

86
migrations/public/versions/9429f244f1a5_moved_storage_and_embedding_licensing_.py Normal file
View File

@@ -0,0 +1,86 @@
"""Moved storage and embedding licensing to Mb iso tokens
Revision ID: 9429f244f1a5
Revises: 48714f1baac5
Create Date: 2024-10-04 08:07:47.976861
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = '9429f244f1a5'
down_revision = '48714f1baac5'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.add_column(sa.Column('max_storage_mb', sa.Integer(), nullable=False))
batch_op.add_column(sa.Column('additional_storage_price', sa.Float(), nullable=False))
batch_op.add_column(sa.Column('included_embedding_mb', sa.Integer(), nullable=False))
batch_op.add_column(sa.Column('additional_embedding_price', sa.Numeric(precision=10, scale=4), nullable=False))
batch_op.add_column(sa.Column('overage_embedding', sa.Float(), nullable=False))
batch_op.add_column(sa.Column('overage_interaction', sa.Float(), nullable=False))
batch_op.drop_column('additional_storage_token_price')
batch_op.drop_column('additional_embedding_token_price')
batch_op.drop_column('max_storage_tokens')
batch_op.drop_column('allow_overage')
batch_op.drop_column('included_embedding_tokens')
with op.batch_alter_table('license_tier', schema=None) as batch_op:
batch_op.add_column(sa.Column('max_storage_mb', sa.Integer(), nullable=False))
batch_op.add_column(sa.Column('additional_storage_price_d', sa.Numeric(precision=10, scale=4), nullable=False))
batch_op.add_column(sa.Column('additional_storage_price_e', sa.Numeric(precision=10, scale=4), nullable=False))
batch_op.add_column(sa.Column('included_embedding_mb', sa.Integer(), nullable=False))
batch_op.add_column(sa.Column('additional_embedding_price_d', sa.Numeric(precision=10, scale=4), nullable=False))
batch_op.add_column(sa.Column('additional_embedding_price_e', sa.Numeric(precision=10, scale=4), nullable=False))
batch_op.add_column(sa.Column('standard_overage_embedding', sa.Float(), nullable=False))
batch_op.add_column(sa.Column('standard_overage_interaction', sa.Float(), nullable=False))
batch_op.drop_column('max_storage_tokens')
batch_op.drop_column('additional_storage_token_price_e')
batch_op.drop_column('allow_overage')
batch_op.drop_column('additional_embedding_token_price_d')
batch_op.drop_column('included_embedding_tokens')
batch_op.drop_column('additional_embedding_token_price_e')
batch_op.drop_column('additional_storage_token_price_d')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license_tier', schema=None) as batch_op:
batch_op.add_column(sa.Column('additional_storage_token_price_d', sa.NUMERIC(precision=10, scale=4), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('additional_embedding_token_price_e', sa.NUMERIC(precision=10, scale=4), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('included_embedding_tokens', sa.INTEGER(), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('additional_embedding_token_price_d', sa.NUMERIC(precision=10, scale=4), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('allow_overage', sa.BOOLEAN(), autoincrement=False, nullable=True))
batch_op.add_column(sa.Column('additional_storage_token_price_e', sa.NUMERIC(precision=10, scale=4), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('max_storage_tokens', sa.INTEGER(), autoincrement=False, nullable=False))
batch_op.drop_column('standard_overage_interaction')
batch_op.drop_column('standard_overage_embedding')
batch_op.drop_column('additional_embedding_price_e')
batch_op.drop_column('additional_embedding_price_d')
batch_op.drop_column('included_embedding_mb')
batch_op.drop_column('additional_storage_price_e')
batch_op.drop_column('additional_storage_price_d')
batch_op.drop_column('max_storage_mb')
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.add_column(sa.Column('included_embedding_tokens', sa.INTEGER(), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('allow_overage', sa.BOOLEAN(), autoincrement=False, nullable=True))
batch_op.add_column(sa.Column('max_storage_tokens', sa.INTEGER(), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('additional_embedding_token_price', sa.NUMERIC(precision=10, scale=4), autoincrement=False, nullable=False))
batch_op.add_column(sa.Column('additional_storage_token_price', sa.DOUBLE_PRECISION(precision=53), autoincrement=False, nullable=False))
batch_op.drop_column('overage_interaction')
batch_op.drop_column('overage_embedding')
batch_op.drop_column('additional_embedding_price')
batch_op.drop_column('included_embedding_mb')
batch_op.drop_column('additional_storage_price')
batch_op.drop_column('max_storage_mb')
# ### end Alembic commands ###

58
migrations/public/versions/d616ea937a6a_corrections_on_licensing_models.py Normal file
View File

@@ -0,0 +1,58 @@
"""Corrections on Licensing models
Revision ID: d616ea937a6a
Revises: 560d08d91e5b
Create Date: 2024-10-02 15:39:27.668741
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = 'd616ea937a6a'
down_revision = '560d08d91e5b'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('business_event_log', schema=None) as batch_op:
batch_op.drop_constraint('business_event_log_license_usage_id_fkey', type_='foreignkey')
batch_op.create_foreign_key(None, 'license_usage', ['license_usage_id'], ['id'], referent_schema='public')
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.drop_constraint('license_tenant_id_fkey', type_='foreignkey')
batch_op.drop_constraint('license_tier_id_fkey', type_='foreignkey')
batch_op.create_foreign_key(None, 'license_tier', ['tier_id'], ['id'], referent_schema='public')
batch_op.create_foreign_key(None, 'tenant', ['tenant_id'], ['id'], referent_schema='public')
with op.batch_alter_table('license_usage', schema=None) as batch_op:
batch_op.drop_constraint('license_usage_license_id_fkey', type_='foreignkey')
batch_op.drop_constraint('license_usage_tenant_id_fkey', type_='foreignkey')
batch_op.create_foreign_key(None, 'tenant', ['tenant_id'], ['id'], referent_schema='public')
batch_op.create_foreign_key(None, 'license', ['license_id'], ['id'], referent_schema='public')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license_usage', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.create_foreign_key('license_usage_tenant_id_fkey', 'tenant', ['tenant_id'], ['id'])
batch_op.create_foreign_key('license_usage_license_id_fkey', 'license', ['license_id'], ['id'])
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.create_foreign_key('license_tier_id_fkey', 'license_tier', ['tier_id'], ['id'])
batch_op.create_foreign_key('license_tenant_id_fkey', 'tenant', ['tenant_id'], ['id'])
with op.batch_alter_table('business_event_log', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.create_foreign_key('business_event_log_license_usage_id_fkey', 'license_usage', ['license_usage_id'], ['id'])
# ### end Alembic commands ###

66
migrations/public/versions/f201bfd23152_add_buckets_to_license_information.py Normal file
View File

@@ -0,0 +1,66 @@
"""Add buckets to License information
Revision ID: f201bfd23152
Revises: d616ea937a6a
Create Date: 2024-10-03 09:44:32.867470
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = 'f201bfd23152'
down_revision = 'd616ea937a6a'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('business_event_log', schema=None) as batch_op:
batch_op.drop_constraint('business_event_log_license_usage_id_fkey', type_='foreignkey')
batch_op.create_foreign_key(None, 'license_usage', ['license_usage_id'], ['id'], referent_schema='public')
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.add_column(sa.Column('additional_storage_bucket', sa.Integer(), nullable=False))
batch_op.drop_constraint('license_tier_id_fkey', type_='foreignkey')
batch_op.drop_constraint('license_tenant_id_fkey', type_='foreignkey')
batch_op.create_foreign_key(None, 'license_tier', ['tier_id'], ['id'], referent_schema='public')
batch_op.create_foreign_key(None, 'tenant', ['tenant_id'], ['id'], referent_schema='public')
with op.batch_alter_table('license_tier', schema=None) as batch_op:
batch_op.add_column(sa.Column('additional_storage_bucket', sa.Integer(), nullable=False))
with op.batch_alter_table('license_usage', schema=None) as batch_op:
batch_op.drop_constraint('license_usage_tenant_id_fkey', type_='foreignkey')
batch_op.drop_constraint('license_usage_license_id_fkey', type_='foreignkey')
batch_op.create_foreign_key(None, 'license', ['license_id'], ['id'], referent_schema='public')
batch_op.create_foreign_key(None, 'tenant', ['tenant_id'], ['id'], referent_schema='public')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('license_usage', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.create_foreign_key('license_usage_license_id_fkey', 'license', ['license_id'], ['id'])
batch_op.create_foreign_key('license_usage_tenant_id_fkey', 'tenant', ['tenant_id'], ['id'])
with op.batch_alter_table('license_tier', schema=None) as batch_op:
batch_op.drop_column('additional_storage_bucket')
with op.batch_alter_table('license', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.create_foreign_key('license_tenant_id_fkey', 'tenant', ['tenant_id'], ['id'])
batch_op.create_foreign_key('license_tier_id_fkey', 'license_tier', ['tier_id'], ['id'])
batch_op.drop_column('additional_storage_bucket')
with op.batch_alter_table('business_event_log', schema=None) as batch_op:
batch_op.drop_constraint(None, type_='foreignkey')
batch_op.create_foreign_key('business_event_log_license_usage_id_fkey', 'license_usage', ['license_usage_id'], ['id'])
# ### end Alembic commands ###

94
migrations/tenant/versions/322d3cf1f17b_documentversion_update_to_bucket_name_.py Normal file
View File

@@ -0,0 +1,94 @@
"""DocumentVersion update to bucket_name, object_name and file_size to better reflet Minio reality
Revision ID: 322d3cf1f17b
Revises: 711a09a77680
Create Date: 2024-10-07 07:45:19.014017
"""
from alembic import op
import sqlalchemy as sa
import pgvector
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import Session
from sqlalchemy import text
from flask import current_app
# revision identifiers, used by Alembic.
revision = '322d3cf1f17b'
down_revision = '711a09a77680'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - Add new fields ###
op.add_column('document_version', sa.Column('bucket_name', sa.String(length=255), nullable=True))
op.add_column('document_version', sa.Column('object_name', sa.String(length=200), nullable=True))
op.add_column('document_version', sa.Column('file_size', sa.Float(), nullable=True))
# ### Upgrade values for bucket_name, object_name and file_size to reflect minio reality ###
from common.models.document import DocumentVersion
from common.extensions import minio_client
from minio.error import S3Error
# Create a connection
connection = op.get_bind()
session = Session(bind=connection)
# Get the current schema name (which should be the tenant ID)
current_schema = connection.execute(text("SELECT current_schema()")).scalar()
tenant_id = int(current_schema)
doc_versions = session.query(DocumentVersion).all()
for doc_version in doc_versions:
try:
object_name = minio_client.generate_object_name(doc_version.doc_id,
doc_version.language,
doc_version.id,
doc_version.file_name)
bucket_name = minio_client.generate_bucket_name(tenant_id)
doc_version.object_name = object_name
doc_version.bucket_name = bucket_name
try:
stat = minio_client.client.stat_object(
bucket_name=bucket_name,
object_name=object_name
)
doc_version.file_size = stat.size / 1048576
current_app.logger.info(f"Processed Upgrade for DocumentVersion {doc_version.id} for Tenant {tenant_id}")
except S3Error as e:
if e.code == "NoSuchKey":
current_app.logger.warning(
f"Object {doc_version.file_location} not found in bucket {doc_version.bucket_name}. Skipping.")
continue # Move to the next item
else:
raise e # Handle other types of S3 errors
except Exception as e:
session.rollback()
current_app.logger.error(f"Couldn't process upgrade for DocumentVersion {doc_version.id} for "
f"Tenant {tenant_id}. Error: {str(e)}")
try:
session.commit()
current_app.logger.info(f"Successfully updated file sizes for tenant schema {current_schema}")
except Exception as e:
session.rollback()
current_app.logger.error(f"Error committing changes for tenant schema {current_schema}: {str(e)}")
# ### commands auto generated by Alembic - Remove old fields ###
# op.drop_column('document_version', 'file_location')
# op.drop_column('document_version', 'file_name')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('document_version', sa.Column('file_name', sa.VARCHAR(length=200), autoincrement=False, nullable=True))
op.add_column('document_version', sa.Column('file_location', sa.VARCHAR(length=255), autoincrement=False, nullable=True))
op.drop_column('document_version', 'file_size')
op.drop_column('document_version', 'object_name')
op.drop_column('document_version', 'bucket_name')
# ### end Alembic commands ###