Token Usage Model

Database model for tracking LLM token consumption and cost estimation.

  • TokenUsage - Tracks token usage per LLM API call with cost estimation

Key Features

  • Token counting for prompt, completion, and total tokens

  • Model and provider tracking

  • Cost estimation based on current pricing

  • Relationships to users, pipelines, and tasks

  • Timestamp tracking for usage analytics

Usage Patterns

# Query total tokens by user
from sqlalchemy import func
from ai4drpm.db.models.token_usage import TokenUsage

total_tokens = session.query(func.sum(TokenUsage.total_tokens)).filter(
    TokenUsage.user_id == user_id
).scalar()

# Query daily usage
from datetime import datetime, timedelta

start_of_day = datetime.now() - timedelta(days=1)
daily_usage = session.query(TokenUsage).filter(
    TokenUsage.user_id == user_id,
    TokenUsage.created_at >= start_of_day
).all()

# Estimate cost
from ai4drpm.db.models.token_usage import estimate_cost

cost = estimate_cost(
    model="mistral-large-latest",
    prompt_tokens=1000,
    completion_tokens=500
)