bifroest/app/migrations/versions/0008_nutzung_und_preise.py
2026-10-27 08:34:17 +01:00

66 lines
3.5 KiB
Python

"""Nutzung je Modellaufruf und Modellpreise
Revision ID: 0008
Revises: 0007
Erstellt: 2026-10-22 09:41:27.635420
"""
import sqlalchemy as sa
from alembic import op
revision = '0008'
down_revision = '0007'
branch_labels = None
depends_on = None
def upgrade() -> None:
op.create_table('usage_records',
sa.Column('id', sa.BigInteger(), sa.Identity(always=False), nullable=False),
sa.Column('time', sa.DateTime(timezone=True), nullable=False),
sa.Column('run_id', sa.Uuid(), nullable=True),
sa.Column('team_id', sa.Integer(), nullable=True),
sa.Column('workspace_id', sa.Integer(), nullable=True),
sa.Column('user_id', sa.Integer(), nullable=True),
sa.Column('model', sa.String(length=200), nullable=False),
sa.Column('input_tokens', sa.Integer(), nullable=False),
sa.Column('cached_tokens', sa.Integer(), nullable=False),
sa.Column('output_tokens', sa.Integer(), nullable=False),
sa.Column('reasoning_tokens', sa.Integer(), nullable=False),
sa.Column('cost', sa.Numeric(precision=14, scale=6), nullable=True),
sa.ForeignKeyConstraint(['team_id'], ['teams.id'], name=op.f('fk_usage_records_team_id_teams'), ondelete='SET NULL'),
sa.ForeignKeyConstraint(['user_id'], ['users.id'], name=op.f('fk_usage_records_user_id_users'), ondelete='SET NULL'),
sa.ForeignKeyConstraint(['workspace_id'], ['workspaces.id'], name=op.f('fk_usage_records_workspace_id_workspaces'), ondelete='SET NULL'),
sa.PrimaryKeyConstraint('id', name=op.f('pk_usage_records'))
)
op.create_index(op.f('ix_usage_records_run_id'), 'usage_records', ['run_id'], unique=False)
op.create_index(op.f('ix_usage_records_team_id'), 'usage_records', ['team_id'], unique=False)
op.create_index(op.f('ix_usage_records_time'), 'usage_records', ['time'], unique=False)
op.create_index(op.f('ix_usage_records_workspace_id'), 'usage_records', ['workspace_id'], unique=False)
op.add_column('models', sa.Column('input_price', sa.Numeric(precision=12, scale=4), nullable=True))
op.add_column('models', sa.Column('cached_price', sa.Numeric(precision=12, scale=4), nullable=True))
op.add_column('models', sa.Column('output_price', sa.Numeric(precision=12, scale=4), nullable=True))
# Bisherigen Verbrauch aus den Antworten übernehmen (je Antwort summiert, ohne Preis)
op.execute("""
INSERT INTO usage_records (time, run_id, team_id, workspace_id, user_id, model, input_tokens, cached_tokens,
output_tokens, reasoning_tokens)
SELECT m.created, m.run_id, w.team_id, c.workspace_id, r.created_by, m.model,
coalesce((m.usage->>'input_tokens')::int, 0), 0, coalesce((m.usage->>'output_tokens')::int, 0), 0
FROM messages m
JOIN chat_sessions c ON c.id = m.chat_session_id
JOIN workspaces w ON w.id = c.workspace_id
LEFT JOIN task_runs r ON r.id = m.run_id
WHERE m.role = 'assistant' AND m.usage IS NOT NULL AND (m.usage->>'total_tokens') IS NOT NULL
""")
def downgrade() -> None:
op.drop_column('models', 'output_price')
op.drop_column('models', 'cached_price')
op.drop_column('models', 'input_price')
op.drop_index(op.f('ix_usage_records_workspace_id'), table_name='usage_records')
op.drop_index(op.f('ix_usage_records_time'), table_name='usage_records')
op.drop_index(op.f('ix_usage_records_team_id'), table_name='usage_records')
op.drop_index(op.f('ix_usage_records_run_id'), table_name='usage_records')
op.drop_table('usage_records')