"""Cost tracking and budget management for the Headroom proxy. Contains the CostTracker class and cost-related helper functions for prefix cache statistics, cost merging, and session summaries. Extracted from server.py for maintainability. """ from __future__ import annotations import logging from collections import deque from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from headroom.proxy.prometheus_metrics import PrometheusMetrics # Try to import LiteLLM for pricing try: import litellm LITELLM_AVAILABLE = True except ImportError: LITELLM_AVAILABLE = False logger = logging.getLogger("headroom.proxy") # Provider-specific cache discount multipliers (what fraction of input price) # Used to calculate dollar savings from prefix caching _CACHE_ECONOMICS = { "anthropic": { "read_multiplier": 0.1, "write_multiplier": 1.25, "label": "Explicit breakpoints, 5-min TTL", }, "openai": { "read_multiplier": 0.5, "write_multiplier": 1.0, "label": "Automatic, no TTL control", }, "gemini": { "read_multiplier": 0.1, "write_multiplier": 1.0, "label": "Explicit cachedContent, configurable TTL", }, "bedrock": { "read_multiplier": 0.1, "write_multiplier": 1.25, "label": "Same as Anthropic (Bedrock)", }, } def _summarize_transforms(transforms: list[str]) -> str: """Collapse repeated transforms into counted summary. e.g. ['router:excluded:tool', 'router:excluded:tool', 'read_lifecycle:stale'] → 'router:excluded:tool*2 read_lifecycle:stale' """ if not transforms: return "none" counts: dict[str, int] = {} for t in transforms: counts[t] = counts.get(t, 0) + 1 parts = [f"{k}*{v}" if v > 1 else k for k, v in counts.items()] return " ".join(parts) def build_prefix_cache_stats( metrics: PrometheusMetrics, cost_tracker: CostTracker | None, ) -> dict: """Build provider-aware prefix cache statistics for the dashboard.""" by_provider = {} totals = { "cache_read_tokens": 0, "cache_write_tokens": 0, "requests": 0, "hit_requests": 0, "bust_count": 0, "bust_write_tokens": 0, "savings_usd": 0.0, "write_premium_usd": 0.0, } for provider, pc in metrics.cache_by_provider.items(): if pc["requests"] == 0: continue econ = _CACHE_ECONOMICS.get(provider, _CACHE_ECONOMICS["anthropic"]) read_mult: float = econ["read_multiplier"] # type: ignore[assignment] write_mult: float = econ["write_multiplier"] # type: ignore[assignment] # Get the base input price per token for the most-used model on this provider input_price_per_token = None if cost_tracker: for model_name in cost_tracker._tokens_sent_by_model: # Match model to provider _openai_prefixes = ("gpt", "o1", "o3", "o4") is_match = ( (provider == "anthropic" and "claude" in model_name) or (provider == "openai" and any(p in model_name for p in _openai_prefixes)) or (provider == "gemini" and "gemini" in model_name) or (provider == "bedrock" and "claude" in model_name) ) if is_match: price_per_1m = cost_tracker._get_list_price(model_name) if price_per_1m: input_price_per_token = price_per_1m / 1_000_000 break # Calculate savings: # Cache reads save (1.0 - read_mult) per token vs uncached input price. # Cache write premium is NOT deducted — it's baseline cost that the # client (e.g. Claude Code) pays regardless of Headroom. We track it # for observability but don't penalise our savings number. read_tokens: int = pc["cache_read_tokens"] # type: ignore[assignment] write_tokens: int = pc["cache_write_tokens"] # type: ignore[assignment] savings_usd = 0.0 write_premium_usd = 0.0 if input_price_per_token: # Savings from reads: tokens * price * (1.0 - read_multiplier) savings_usd = read_tokens * input_price_per_token * (1.0 - read_mult) # Write premium (observability only — not subtracted from savings) if write_mult > 1.0: write_premium_usd = write_tokens * input_price_per_token * (write_mult - 1.0) hit_rate = round(pc["hit_requests"] / pc["requests"] * 100, 1) if pc["requests"] > 0 else 0 provider_stats = { "cache_read_tokens": read_tokens, "cache_write_tokens": write_tokens, "requests": pc["requests"], "hit_requests": pc["hit_requests"], "hit_rate": hit_rate, "bust_count": pc["bust_count"], "bust_write_tokens": pc["bust_write_tokens"], "read_discount": f"{(1.0 - read_mult) * 100:.0f}%", "write_premium": f"{(write_mult - 1.0) * 100:.0f}%" if write_mult > 1.0 else "none", "savings_usd": round(savings_usd, 4), "write_premium_usd": round(write_premium_usd, 4), "net_savings_usd": round(savings_usd, 4), "label": str(econ["label"]), } by_provider[provider] = provider_stats # Accumulate totals totals["cache_read_tokens"] += read_tokens totals["cache_write_tokens"] += write_tokens totals["requests"] += pc["requests"] totals["hit_requests"] += pc["hit_requests"] totals["bust_count"] += pc["bust_count"] totals["bust_write_tokens"] += pc["bust_write_tokens"] totals["savings_usd"] += savings_usd totals["write_premium_usd"] += write_premium_usd totals["net_savings_usd"] = round(totals["savings_usd"], 4) totals["savings_usd"] = round(totals["savings_usd"], 4) totals["write_premium_usd"] = round(totals["write_premium_usd"], 4) totals["hit_rate"] = ( round(totals["hit_requests"] / totals["requests"] * 100, 1) if totals["requests"] > 0 else 0 ) return { "by_provider": by_provider, "totals": totals, "prefix_freeze": { "busts_avoided": metrics.prefix_freeze_busts_avoided, "tokens_preserved": metrics.prefix_freeze_tokens_preserved, "compression_foregone_tokens": metrics.prefix_freeze_compression_foregone, "net_benefit_tokens": ( metrics.prefix_freeze_tokens_preserved - metrics.prefix_freeze_compression_foregone ), }, "attribution": ( "Prefix caching is performed by the LLM provider (Anthropic, OpenAI). " "Headroom reports cache stats as observed from API responses. " "CacheAligner and prefix freeze improve cache hit rates by stabilizing " "the message prefix, but baseline caching happens without Headroom." ), } def merge_cost_stats( cost_stats: dict | None, cache_stats: dict, cli_tokens_avoided: int = 0, ) -> dict | None: """Merge compression, cache, and CLI savings into cost stats. Each savings layer is reported separately with its own scope: - savings_usd: compression savings at model list price (monotonic) - cache_savings_usd: prefix cache discount from provider (separate) - cli_tokens_avoided: tokens filtered by rtk (token count only, no $ estimate) The hero metric (savings_usd) is ONLY compression savings priced at the model's published input rate. Cache and CLI are shown separately. This avoids the non-monotonic moving-average repricing bug (#83). """ if cost_stats is None: return None cache_net = cache_stats.get("totals", {}).get("net_savings_usd", 0.0) compression_savings = cost_stats.get("savings_usd", 0.0) return { **cost_stats, "savings_usd": round(compression_savings, 4), "compression_savings_usd": round(compression_savings, 4), "cache_savings_usd": round(cache_net, 4), "cli_tokens_avoided": cli_tokens_avoided, } def build_session_summary( proxy: Any, metrics: Any, prefix_cache_stats: dict, cli_tokens_avoided: int, total_tokens_before: int, ) -> dict[str, Any]: """Build a human-readable session summary from metrics and request logs. This is the headline view users see first in /stats — designed to answer "is Headroom working?" at a glance. """ # Analyze per-request compression from the logger compressed_requests: list[dict] = [] uncompressed_reasons: dict[str, int] = { "prefix_frozen": 0, "too_small": 0, "passthrough": 0, "no_compressible_content": 0, } if proxy.logger: for entry in proxy.logger._logs: if entry.model and "count_tokens" in entry.model: uncompressed_reasons["passthrough"] += 1 continue if entry.tokens_saved > 0: compressed_requests.append( { "savings_pct": round(entry.savings_percent, 1), "tokens_saved": entry.tokens_saved, "original": entry.input_tokens_original, "optimized": entry.input_tokens_optimized, } ) elif entry.input_tokens_original > 0: # Categorize why it wasn't compressed transforms = entry.transforms_applied or [] if not transforms: # Pipeline returned unchanged — likely all frozen uncompressed_reasons["prefix_frozen"] += 1 elif all("excluded" in t or "protected" in t for t in transforms): uncompressed_reasons["no_compressible_content"] += 1 elif entry.input_tokens_original < 500: uncompressed_reasons["too_small"] += 1 else: uncompressed_reasons["prefix_frozen"] += 1 # Compute compression stats for requests that DID compress avg_compression = 0.0 best_compression = 0.0 best_detail = "" if compressed_requests: avg_compression = round( sum(r["savings_pct"] for r in compressed_requests) / len(compressed_requests), 1, ) best = max(compressed_requests, key=lambda r: r["savings_pct"]) best_compression = best["savings_pct"] best_detail = f"{best['original']:,} → {best['optimized']:,} tokens" # Cost summary — savings_usd is compression savings at model list price (monotonic) cost_stats = proxy.cost_tracker.stats() if proxy.cost_tracker else {} cost_with = cost_stats.get("cost_with_headroom_usd", 0.0) compression_savings = cost_stats.get("savings_usd", 0.0) cache_net = prefix_cache_stats.get("totals", {}).get("net_savings_usd", 0.0) total_saved_usd = round(compression_savings, 2) cost_without = cost_with + compression_savings savings_pct_cost = round(total_saved_usd / cost_without * 100, 1) if cost_without > 0 else 0.0 # Primary models used models = dict(metrics.requests_by_model) primary_model = max(models, key=lambda k: models[k]) if models else "unknown" api_requests = sum(v for k, v in models.items() if "count_tokens" not in k) # Build the summary summary: dict[str, Any] = { "mode": proxy.config.mode, "api_requests": api_requests, "primary_model": primary_model, "compression": { "requests_compressed": len(compressed_requests), "avg_compression_pct": avg_compression, "best_compression_pct": best_compression, "best_detail": best_detail, "total_tokens_removed": metrics.tokens_saved_total, }, "uncompressed_requests": {k: v for k, v in uncompressed_reasons.items() if v > 0}, "cost": { "without_headroom_usd": round(cost_without, 2), "with_headroom_usd": round(cost_with, 2), "total_saved_usd": total_saved_usd, "savings_pct": savings_pct_cost, "breakdown": { "cache_savings_usd": round(cache_net, 2), "compression_savings_usd": round(compression_savings, 2), }, }, } # Add tip if token_headroom mode would help if proxy.config.mode == "cost_savings" and uncompressed_reasons["prefix_frozen"] > 10: summary["tip"] = ( "Most requests are prefix-frozen. Set HEADROOM_MODE=token_headroom " "to compress frozen messages and extend your session by ~25-35%." ) return summary class CostTracker: """Track costs and enforce budgets. Cost history is automatically pruned to prevent unbounded memory growth: - Entries older than 24 hours are removed - Maximum of 100,000 entries are kept Uses LiteLLM's community-maintained pricing database for accurate costs. See: https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json """ MAX_COST_ENTRIES = 100_000 COST_RETENTION_HOURS = 24 def __init__(self, budget_limit_usd: float | None = None, budget_period: str = "daily"): self.budget_limit_usd = budget_limit_usd self.budget_period = budget_period # Cost tracking - using deque for efficient left-side removal self._costs: deque[tuple[datetime, float]] = deque(maxlen=self.MAX_COST_ENTRIES) self._last_prune_time: datetime = datetime.now() # Token savings per model (exact, no dollar estimation) self._tokens_saved_by_model: dict[str, int] = {} self._tokens_sent_by_model: dict[str, int] = {} self._requests_by_model: dict[str, int] = {} # API-reported cache breakdown per model (for accurate cost calculation) self._api_cache_read_by_model: dict[str, int] = {} self._api_cache_write_by_model: dict[str, int] = {} self._api_uncached_by_model: dict[str, int] = {} # Cache resolved model names to avoid repeated litellm lookups. # This is critical: litellm.cost_per_token() is synchronous and can block # the async event loop if it triggers I/O (lazy model info download). _resolved_model_cache: dict[str, str] = {} @classmethod def _resolve_litellm_model(cls, model: str) -> str: """Resolve model name to one LiteLLM recognizes, adding provider prefix if needed. Results are cached per model name to avoid blocking the event loop with repeated synchronous litellm lookups. """ if model in cls._resolved_model_cache: return cls._resolved_model_cache[model] resolved = cls._resolve_litellm_model_uncached(model) cls._resolved_model_cache[model] = resolved return resolved @staticmethod def _resolve_litellm_model_uncached(model: str) -> str: """Uncached resolution — called once per unique model name.""" if not LITELLM_AVAILABLE: return model # Try as-is first try: litellm.cost_per_token(model=model, prompt_tokens=1, completion_tokens=0) return model except Exception: pass # Try with provider prefix prefixes = { "claude-": "anthropic/", "gpt-": "openai/", "o1-": "openai/", "o3-": "openai/", "o4-": "openai/", "gemini-": "google/", } for pattern, prefix in prefixes.items(): if model.startswith(pattern): prefixed = f"{prefix}{model}" try: litellm.cost_per_token(model=prefixed, prompt_tokens=1, completion_tokens=0) return prefixed except Exception: break return model def estimate_cost( self, model: str, input_tokens: int, output_tokens: int, cache_read_tokens: int = 0, cache_write_tokens: int = 0, ) -> float | None: """Estimate cost in USD using LiteLLM's pricing database. LiteLLM natively handles cache_read and cache_creation pricing for all providers (Anthropic, OpenAI, Google, etc.) in a single call. Args: model: Model name for pricing lookup input_tokens: Non-cached input tokens (excludes cache_read) output_tokens: Output tokens cache_read_tokens: Tokens served from cache (~10% of input rate) cache_write_tokens: Tokens written to cache (~125% of input rate) """ if not LITELLM_AVAILABLE: logger.warning("LiteLLM not available - cannot calculate costs") return None try: resolved_model = self._resolve_litellm_model(model) # litellm.cost_per_token handles all token types natively: # prompt_tokens at input rate, cache_read at ~10%, cache_creation at ~125% input_cost, output_cost = litellm.cost_per_token( model=resolved_model, prompt_tokens=input_tokens, completion_tokens=output_tokens, cache_read_input_tokens=cache_read_tokens, cache_creation_input_tokens=cache_write_tokens, ) total_cost = input_cost + output_cost return float(total_cost) if total_cost > 0 else None except Exception as e: logger.warning(f"Failed to get pricing for model {model}: {e}") return None def _prune_old_costs(self): """Remove cost entries older than retention period. Called periodically (every 5 minutes) to prevent unbounded memory growth. The deque maxlen provides a hard cap, but time-based pruning keeps memory usage proportional to actual traffic patterns. """ now = datetime.now() # Only prune every 5 minutes to avoid overhead if (now - self._last_prune_time).total_seconds() < 300: return self._last_prune_time = now cutoff = now - timedelta(hours=self.COST_RETENTION_HOURS) # Remove entries from the left (oldest) while they're older than cutoff while self._costs and self._costs[0][0] < cutoff: self._costs.popleft() def record_tokens( self, model: str, tokens_saved: int, tokens_sent: int, cache_read_tokens: int = 0, cache_write_tokens: int = 0, uncached_tokens: int = 0, ): """Record token counts per model. Args: model: Model name. tokens_saved: Tokens removed by compression (Headroom's count). tokens_sent: Compressed message tokens sent (Headroom's count). cache_read_tokens: Cache read tokens from API response usage. cache_write_tokens: Cache write tokens from API response usage. uncached_tokens: Non-cached input tokens from API response usage. """ self._tokens_saved_by_model[model] = ( self._tokens_saved_by_model.get(model, 0) + tokens_saved ) self._tokens_sent_by_model[model] = self._tokens_sent_by_model.get(model, 0) + tokens_sent self._requests_by_model[model] = self._requests_by_model.get(model, 0) + 1 self._api_cache_read_by_model[model] = ( self._api_cache_read_by_model.get(model, 0) + cache_read_tokens ) self._api_cache_write_by_model[model] = ( self._api_cache_write_by_model.get(model, 0) + cache_write_tokens ) self._api_uncached_by_model[model] = ( self._api_uncached_by_model.get(model, 0) + uncached_tokens ) def get_period_cost(self) -> float: """Get cost for current budget period.""" now = datetime.now() if self.budget_period == "hourly": cutoff = now - timedelta(hours=1) elif self.budget_period == "daily": cutoff = now.replace(hour=0, minute=0, second=0, microsecond=0) else: # monthly cutoff = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) return sum(cost for ts, cost in self._costs if ts >= cutoff) def check_budget(self) -> tuple[bool, float]: """Check if within budget. Returns (allowed, remaining).""" if self.budget_limit_usd is None: return True, float("inf") period_cost = self.get_period_cost() remaining = self.budget_limit_usd - period_cost return remaining > 0, max(0, remaining) def _get_list_price(self, model: str) -> float | None: """Get list input price per 1M tokens for a model.""" if not LITELLM_AVAILABLE: return None try: resolved = self._resolve_litellm_model(model) info = litellm.model_cost.get(resolved, {}) cost_per_token = info.get("input_cost_per_token") return cost_per_token * 1_000_000 if cost_per_token else None except Exception: return None def _get_cache_prices(self, model: str) -> tuple[float, float, float] | None: """Get per-token prices for cache read, cache write, and uncached input. Returns (cache_read, cache_write, uncached) per-token costs, or None if pricing is unavailable. Uses LiteLLM's native cache pricing data. """ if not LITELLM_AVAILABLE: return None try: resolved = self._resolve_litellm_model(model) info = litellm.model_cost.get(resolved, {}) uncached = info.get("input_cost_per_token") if not uncached: return None cache_read = info.get("cache_read_input_token_cost", uncached) cache_write = info.get("cache_creation_input_token_cost", uncached) return (cache_read, cache_write, uncached) except Exception: return None def stats(self) -> dict: """Get token statistics per model.""" per_model = {} total_saved = 0 for model in sorted(self._tokens_saved_by_model.keys()): saved = self._tokens_saved_by_model[model] sent = self._tokens_sent_by_model.get(model, 0) reqs = self._requests_by_model.get(model, 0) total_saved += saved per_model[model] = { "requests": reqs, "tokens_saved": saved, "tokens_sent": sent, "reduction_pct": round(saved / (saved + sent) * 100, 1) if (saved + sent) > 0 else 0, } # Compute actual input cost using API-reported cache breakdown and # LiteLLM's per-category pricing (cache reads discounted, writes at # premium, uncached at list). Falls back to list price when cache # data is unavailable. cost_with_headroom = 0.0 total_billed_input_tokens = 0 total_input_tokens = 0 for model in self._tokens_saved_by_model: saved = self._tokens_saved_by_model[model] sent = self._tokens_sent_by_model.get(model, 0) cr = self._api_cache_read_by_model.get(model, 0) cw = self._api_cache_write_by_model.get(model, 0) uncached = self._api_uncached_by_model.get(model, 0) total_input_tokens += sent prices = self._get_cache_prices(model) if prices: cr_price, cw_price, uncached_price = prices if cr + cw + uncached > 0: # Use API's real cache breakdown with LiteLLM pricing model_cost = cr * cr_price + cw * cw_price + uncached * uncached_price billed_tokens = cr + cw + uncached else: # No cache data from API — fall back to list price model_cost = sent * uncached_price billed_tokens = sent cost_with_headroom += model_cost total_billed_input_tokens += billed_tokens # Compression savings: price saved tokens at the model's list input price. # This is simple, monotonic, and transparent — each saved token is valued # at the published $/token rate for its model. Not affected by cache mix. savings_usd = 0.0 for model in self._tokens_saved_by_model: saved = self._tokens_saved_by_model[model] if saved <= 0: continue prices = self._get_cache_prices(model) if prices: _cr_price, _cw_price, uncached_price = prices savings_usd += saved * uncached_price return { "total_tokens_saved": total_saved, "total_input_tokens": total_input_tokens, "total_input_cost_usd": round(cost_with_headroom, 4), "per_model": per_model, "cost_with_headroom_usd": round(cost_with_headroom, 4), "savings_usd": round(savings_usd, 4), }