Cost Management › Cost Allocation

# Cost allocation system implementation
import time
from collections import defaultdict
from dataclasses import dataclass


@dataclass
class RequestRecord:
    """A single request record for cost attribution."""

    request_id: str
    user_id: str
    feature: str
    team: str
    timestamp: float
    prompt_tokens: int
    completion_tokens: int
    model_tier: str  # "small", "medium", "large"


# Cost rates by model tier (USD per 1K tokens)
TIER_COSTS = {
    "small": {"input": 0.0002, "output": 0.0005},  # 7B model
    "medium": {"input": 0.0010, "output": 0.0020},  # 13B model
    "large": {"input": 0.0015, "output": 0.0050},  # 70B model
}


def compute_request_cost(record: RequestRecord) -> float:
    """Compute the dollar cost of a single request."""
    rates = TIER_COSTS[record.model_tier]
    input_cost = (record.prompt_tokens / 1000) * rates["input"]
    output_cost = (record.completion_tokens / 1000) * rates["output"]
    return input_cost + output_cost


def aggregate_costs(records: list) -> dict:
    """Aggregate costs by user, feature, and team."""
    by_user = defaultdict(float)
    by_feature = defaultdict(float)
    by_team = defaultdict(float)
    by_tier = defaultdict(float)

    for record in records:
        cost = compute_request_cost(record)
        by_user[record.user_id] += cost
        by_feature[record.feature] += cost
        by_team[record.team] += cost
        by_tier[record.model_tier] += cost

    return {
        "by_user": dict(by_user),
        "by_feature": dict(by_feature),
        "by_team": dict(by_team),
        "by_tier": dict(by_tier),
        "total": sum(by_user.values()),
    }


# Simulate request records
import random

random.seed(42)

users = [f"user_{i:04d}" for i in range(200)]
features = ["autocomplete", "summarization", "chat", "search", "translation"]
teams = ["product_a", "product_b", "infra", "research"]
tiers = ["small", "medium", "large"]
tier_weights = [0.6, 0.3, 0.1]  # Most requests use smaller, cheaper models

simulated_records = []
for i in range(5000):
    prompt_len = max(50, int(random.lognormvariate(6, 0.8)))
    completion_len = max(20, int(random.lognormvariate(5, 0.7)))
    tier = random.choices(tiers, weights=tier_weights)[0]
    simulated_records.append(
        RequestRecord(
            request_id=f"req_{i:06d}",
            user_id=random.choice(users),
            feature=random.choice(features),
            team=random.choice(teams),
            timestamp=time.time() - random.uniform(0, 86400),
            prompt_tokens=min(prompt_len, 4096),
            completion_tokens=min(completion_len, 1024),
            model_tier=tier,
        )
    )

allocation = aggregate_costs(simulated_records)

# Per-user cost limiting with token budgets
class UserCostTracker:
    """Track and enforce per-user cost limits."""

    def __init__(self, daily_limit_usd: float):
        self.daily_limit_usd = daily_limit_usd
        self.user_costs = defaultdict(float)  # user_id -> daily cost so far
        self.user_request_counts = defaultdict(int)

    def check_and_deduct(self, user_id: str, estimated_cost: float) -> tuple:
        """
        Check if user is within budget and deduct the estimated cost.
        Returns (allowed: bool, remaining_budget: float).
        """
        current_spend = self.user_costs[user_id]
        remaining = self.daily_limit_usd - current_spend

        if estimated_cost > remaining:
            return False, remaining

        self.user_costs[user_id] += estimated_cost
        self.user_request_counts[user_id] += 1
        return True, remaining - estimated_cost

    def get_usage_report(self) -> dict:
        """Return usage stats for all users who made requests."""
        return {
            user_id: {
                "total_cost": self.user_costs[user_id],
                "requests": self.user_request_counts[user_id],
                "utilization_pct": self.user_costs[user_id]
                / self.daily_limit_usd
                * 100,
            }
            for user_id in self.user_costs
        }


# Simulate per-user tracking with $0.50 daily limit
tracker = UserCostTracker(daily_limit_usd=0.50)
blocked_count = 0
processed_count = 0

for record in simulated_records[:500]:
    cost = compute_request_cost(record)
    allowed, remaining = tracker.check_and_deduct(record.user_id, cost)
    if allowed:
        processed_count += 1
    else:
        blocked_count += 1

usage = tracker.get_usage_report()