LLM Cost Attribution and Chargeback for Multi-Tenant Teams

import time
from dataclasses import dataclass
from anthropic import Anthropic
 
client = Anthropic()
 
@dataclass
class UsageContext:
    team: str
    feature: str
    customer_id: str | None = None
    environment: str = "production"
 
def tracked_call(messages: list, model: str, max_tokens: int, context: UsageContext) -> dict:
    start = time.time()
    
    response = client.messages.create(
        model=model,
        max_tokens=max_tokens,
        messages=messages
    )
    
    usage_record = {
        "timestamp": time.time(),
        "team": context.team,
        "feature": context.feature,
        "customer_id": context.customer_id,
        "environment": context.environment,
        "model": model,
        "input_tokens": response.usage.input_tokens,
        "output_tokens": response.usage.output_tokens,
        "latency_seconds": time.time() - start,
    }
    
    log_usage(usage_record)
    return response

# Pricing changes — keep this centralized and easy to update
MODEL_PRICING = {
    "claude-sonnet-4-6": {"input": 3.00, "output": 15.00},   # per million tokens
    "claude-haiku-4-5-20251001": {"input": 0.80, "output": 4.00},
    "claude-opus-4-8": {"input": 15.00, "output": 75.00},
}
 
def calculate_cost(record: dict) -> float:
    pricing = MODEL_PRICING[record["model"]]
    input_cost = (record["input_tokens"] / 1_000_000) * pricing["input"]
    output_cost = (record["output_tokens"] / 1_000_000) * pricing["output"]
    return round(input_cost + output_cost, 6)

import sqlite3  # use Postgres/Timescale in production; this is the simplest version
 
def log_usage(record: dict):
    cost = calculate_cost(record)
    
    conn = sqlite3.connect("llm_usage.db")
    conn.execute("""
        INSERT INTO llm_usage 
        (timestamp, team, feature, customer_id, environment, model, 
         input_tokens, output_tokens, cost_usd)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        record["timestamp"], record["team"], record["feature"], 
        record["customer_id"], record["environment"], record["model"],
        record["input_tokens"], record["output_tokens"], cost
    ))
    conn.commit()

CREATE TABLE llm_usage (
    timestamp REAL,
    team TEXT,
    feature TEXT,
    customer_id TEXT,
    environment TEXT,
    model TEXT,
    input_tokens INTEGER,
    output_tokens INTEGER,
    cost_usd REAL
);
 
CREATE INDEX idx_usage_team_time ON llm_usage(team, timestamp);
CREATE INDEX idx_usage_customer ON llm_usage(customer_id);

def monthly_chargeback_report(year: int, month: int) -> dict:
    conn = sqlite3.connect("llm_usage.db")
    
    rows = conn.execute("""
        SELECT team, feature, 
               SUM(cost_usd) as total_cost,
               SUM(input_tokens + output_tokens) as total_tokens,
               COUNT(*) as request_count
        FROM llm_usage
        WHERE strftime('%Y', datetime(timestamp, 'unixepoch')) = ?
          AND strftime('%m', datetime(timestamp, 'unixepoch')) = ?
        GROUP BY team, feature
        ORDER BY total_cost DESC
    """, (str(year), f"{month:02d}")).fetchall()
    
    return [
        {"team": r[0], "feature": r[1], "cost_usd": round(r[2], 2), 
         "tokens": r[3], "requests": r[4]}
        for r in rows
    ]

def customer_usage_for_billing(customer_id: str, start_date: str, end_date: str) -> dict:
    conn = sqlite3.connect("llm_usage.db")
    
    result = conn.execute("""
        SELECT SUM(cost_usd), SUM(input_tokens + output_tokens), COUNT(*)
        FROM llm_usage
        WHERE customer_id = ? AND timestamp BETWEEN ? AND ?
    """, (customer_id, start_date, end_date)).fetchone()
    
    return {
        "customer_id": customer_id,
        "raw_llm_cost": round(result[0] or 0, 2),
        "total_tokens": result[1] or 0,
        "request_count": result[2] or 0
    }

def check_team_budget_alerts(daily_budgets: dict[str, float]):
    """Run hourly — catch a team blowing through their daily budget before month-end surprise."""
    conn = sqlite3.connect("llm_usage.db")
    today_start = get_today_start_timestamp()
    
    for team, budget in daily_budgets.items():
        spent_today = conn.execute("""
            SELECT SUM(cost_usd) FROM llm_usage 
            WHERE team = ? AND timestamp >= ?
        """, (team, today_start)).fetchone()[0] or 0
        
        if spent_today > budget * 0.8:
            send_slack_alert(
                f"⚠️ {team} has spent ${spent_today:.2f} of ${budget:.2f} "
                f"daily LLM budget ({spent_today/budget:.0%})"
            )