diff --git a/headroom/pricing/litellm_pricing.py b/headroom/pricing/litellm_pricing.py index 9a4252689..0b97e4e58 100644 --- a/headroom/pricing/litellm_pricing.py +++ b/headroom/pricing/litellm_pricing.py @@ -282,6 +282,47 @@ def estimate_cost( return input_cost + output_cost +def estimate_cost_from_tokens( + model: str, + input_tokens: int = 0, + output_tokens: int = 0, + cached_tokens: int = 0, +) -> float | None: + """Cost for one request from token counts, using LiteLLM's own cost model. + + Prefer this over :func:`estimate_cost` whenever a request may carry cached + tokens or exceed a model's long-context threshold. Flat per-1M rates cannot + express either: cache reads bill at their own rate, and on Anthropic's + Sonnet 4 / 4.5 family a prompt over 200K re-prices the *whole* request -- + input, output and cache alike. ``litellm.cost_per_token`` applies both. + + ``input_tokens`` is the TOTAL prompt, ``cached_tokens`` included. LiteLLM + subtracts the cached portion itself and tests the long-context threshold + against the total, so passing a cache-exclusive count would both + double-discount the cached tokens and understate the threshold. + + Returns ``None`` when LiteLLM is unavailable (the dependency is gated + ``python_version < '3.14'``) or doesn't know the model -- the caller's cue + to fall back to its own table. + """ + if not LITELLM_AVAILABLE: + return None + candidate = next((c for c in pricing_lookup_candidates(model) if c in litellm.model_cost), None) + if candidate is None: + return None + try: + prompt_cost, completion_cost = litellm.cost_per_token( + model=candidate, + prompt_tokens=input_tokens, + completion_tokens=output_tokens, + cache_read_input_tokens=cached_tokens, + ) + except Exception as exc: # pragma: no cover - depends on litellm internals + logger.debug("litellm.cost_per_token failed for %s: %s", candidate, exc) + return None + return float(prompt_cost) + float(completion_cost) + + def list_available_models() -> list[str]: """List all models with pricing info in LiteLLM's database. diff --git a/headroom/providers/anthropic.py b/headroom/providers/anthropic.py index b16251621..8d4b8313e 100644 --- a/headroom/providers/anthropic.py +++ b/headroom/providers/anthropic.py @@ -24,6 +24,7 @@ import warnings from typing import Any, cast from headroom import paths as _paths +from headroom.pricing.litellm_pricing import estimate_cost_from_tokens from headroom.tokenizers.base import ( TokenCountCache, coerce_countable_text, @@ -67,6 +68,21 @@ def sanitize_anthropic_model_id(model: str) -> str: return _DANGLING_ANSI_STYLE_SUFFIX_RE.sub("", cleaned) +# `[1m]` is not only an ANSI artifact: Claude Code appends it to a model id to +# request the 1M context tier, and only then sends the `context-1m` beta header +# (#1158). Upstream rejects the suffix, so `sanitize_anthropic_model_id()` must +# keep stripping it before forwarding (#2027) — but the tier it encodes has to +# be read off the id *before* that happens, or a 1M request gets budgeted as if +# it were the base model's window. +_CONTEXT_1M_SUFFIX_RE = re.compile(r"(?:\[1m\])+$") +CONTEXT_1M_TOKENS = 1_000_000 + + +def has_context_1m_suffix(model: str) -> bool: + """Return True if ``model`` carries Claude Code's ``[1m]`` 1M-tier marker.""" + return bool(_CONTEXT_1M_SUFFIX_RE.search(_ANSI_ESCAPE_RE.sub("", str(model)).strip())) + + def sanitize_anthropic_model_metadata(value: Any) -> Any: """Strip model-id styling artifacts from Anthropic model metadata payloads.""" if isinstance(value, list): @@ -154,6 +170,40 @@ ANTHROPIC_PRICING: dict[str, dict[str, float]] = { "claude-3-haiku-20240307": {"input": 0.25, "output": 1.25, "cached_input": 0.03}, } +# Anthropic's long-context premium. On models that reach 1M over a 200K base, +# a prompt above 200K re-prices the *entire* request -- input, output and cache +# alike -- rather than only the tokens past the threshold. Multipliers are +# derived from LiteLLM's `*_above_200k_tokens` fields ($3->$6 in, $15->$22.50 +# out, $0.30->$0.60 cache read). +# +# Only the Sonnet 4 / 4.5 family is tiered: Opus, and Sonnet 4.6 onward, are +# flat-rated across their whole window. This is the same population that needs +# the `[1m]` suffix to reach 1M at all, so a session that fills the window this +# unlocks is billed at these rates. +_LONG_CONTEXT_THRESHOLD = 200_000 +_LONG_CONTEXT_PREMIUM: dict[str, float] = {"input": 2.0, "output": 1.5, "cached_input": 2.0} +_LONG_CONTEXT_TIERED_MODELS = ( + "claude-sonnet-4-5", + "claude-sonnet-4-20250514", + "claude-4-sonnet-20250514", +) + + +def _apply_long_context_premium( + model: str, pricing: dict[str, float], input_tokens: int +) -> dict[str, float]: + """Return ``pricing`` scaled by the long-context premium where it applies. + + Used only on the manual fallback path; the LiteLLM path already applies the + published above-threshold rates itself. + """ + if input_tokens <= _LONG_CONTEXT_THRESHOLD: + return pricing + if not any(model.startswith(tiered) for tiered in _LONG_CONTEXT_TIERED_MODELS): + return pricing + return {key: rate * _LONG_CONTEXT_PREMIUM.get(key, 1.0) for key, rate in pricing.items()} + + # Default limits for pattern-based inference # Used when a model isn't in the explicit list but matches a known pattern _PATTERN_DEFAULTS = { @@ -605,8 +655,16 @@ class AnthropicProvider(Provider): 6. Pattern-based inference (opus/sonnet/haiku) 7. Default fallback (200K for any Claude model) + A ``[1m]`` suffix raises the result to at least 1M: the caller asked for + the 1M tier and Claude Code sent the `context-1m` beta header, so the + real upstream window is 1M even when the base model's default is 200K. + Never raises an exception - uses sensible defaults for unknown models. """ + if has_context_1m_suffix(model): + # Recursion terminates: the sanitized id has no `[1m]` left. + base = self.get_context_limit(sanitize_anthropic_model_id(model)) + return max(base, CONTEXT_1M_TOKENS) model = sanitize_anthropic_model_id(model) # Check explicit and loaded limits if model in self._context_limits: @@ -685,58 +743,38 @@ class AnthropicProvider(Provider): """Estimate cost for a request. Tries LiteLLM first for up-to-date pricing, falls back to manual pricing. + Both paths apply Anthropic's long-context premium: on the Sonnet 4 / 4.5 + family a prompt over 200K re-prices the whole request (see + ``_LONG_CONTEXT_PREMIUM``). """ model = sanitize_anthropic_model_id(model) - # Try LiteLLM first for cost estimation - litellm, litellm_get_model_info = _get_litellm_clients() - if litellm is not None: - try: - cost = litellm.completion_cost( - model=model, - prompt="", - completion="", - prompt_tokens=input_tokens - cached_tokens, - completion_tokens=output_tokens, - ) - # Add cached token cost if applicable - if cached_tokens > 0: - try: - # Get cached input pricing from LiteLLM model info - info = ( - litellm_get_model_info(model) - if litellm_get_model_info is not None - else None - ) - if info and "input_cost_per_token" in info: - # LiteLLM typically applies 90% discount for cached tokens - cached_cost = cached_tokens * info["input_cost_per_token"] * 0.1 - cost += cached_cost - except Exception: - # Fall back to manual cached pricing - pricing = self._get_pricing(model) - if pricing: - cached_cost = (cached_tokens / 1_000_000) * pricing.get( - "cached_input", pricing["input"] - ) - cost += cached_cost - return cost # type: ignore[no-any-return] - except Exception as e: - logger.debug(f"LiteLLM cost estimation failed for {model}: {e}") + # LiteLLM knows per-model cache and long-context rates, so let it price + # the whole request rather than rebuilding the rate card here. + cost = estimate_cost_from_tokens( + model, + input_tokens=input_tokens, + output_tokens=output_tokens, + cached_tokens=cached_tokens, + ) + if cost is not None: + return cost # Fall back to manual pricing pricing = self._get_pricing(model) if not pricing: return None + rates = _apply_long_context_premium(model, pricing, input_tokens) + # Calculate cost non_cached_input = input_tokens - cached_tokens cost = ( - (non_cached_input / 1_000_000) * pricing["input"] - + (cached_tokens / 1_000_000) * pricing.get("cached_input", pricing["input"]) - + (output_tokens / 1_000_000) * pricing["output"] + (non_cached_input / 1_000_000) * rates["input"] + + (cached_tokens / 1_000_000) * rates.get("cached_input", rates["input"]) + + (output_tokens / 1_000_000) * rates["output"] ) - return cost # type: ignore[no-any-return] + return cost def _get_pricing(self, model: str) -> dict[str, float] | None: """Get pricing for a model with fallback logic.""" diff --git a/headroom/providers/cohere.py b/headroom/providers/cohere.py index 2b205a0a2..d557cc511 100644 --- a/headroom/providers/cohere.py +++ b/headroom/providers/cohere.py @@ -21,6 +21,7 @@ import warnings from datetime import date from typing import Any +from headroom.pricing.litellm_pricing import estimate_cost_from_tokens from headroom.tokenizers import EstimatingTokenCounter from .base import Provider, TokenCounter @@ -326,18 +327,13 @@ class CohereProvider(Provider): # Try LiteLLM first if LITELLM_AVAILABLE: for model_variant in [f"cohere/{model}", model]: - try: - cost = litellm.completion_cost( - model=model_variant, - prompt="", - completion="", - prompt_tokens=input_tokens, - completion_tokens=output_tokens, - ) - if cost is not None: - return float(cost) - except Exception: - pass + cost = estimate_cost_from_tokens( + model_variant, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + if cost is not None: + return float(cost) # Fallback to built-in pricing model_lower = model.lower() diff --git a/headroom/providers/google.py b/headroom/providers/google.py index 8612d5502..4dc919df3 100644 --- a/headroom/providers/google.py +++ b/headroom/providers/google.py @@ -26,6 +26,7 @@ from datetime import date from typing import Any from headroom.models.registry import ModelRegistry +from headroom.pricing.litellm_pricing import estimate_cost_from_tokens from headroom.tokenizers import EstimatingTokenCounter from .base import Provider, TokenCounter @@ -346,18 +347,13 @@ class GoogleProvider(Provider): model_lower, # gemini-1.5-pro ] for variant in model_variants: - try: - cost = litellm.completion_cost( - model=variant, - prompt="", - completion="", - prompt_tokens=input_tokens, - completion_tokens=output_tokens, - ) - if cost is not None: - return cost - except Exception: - continue + cost = estimate_cost_from_tokens( + variant, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + if cost is not None: + return cost # Fallback to hardcoded pricing input_price, output_price = None, None diff --git a/headroom/providers/litellm.py b/headroom/providers/litellm.py index 95d13311a..e02f3f683 100644 --- a/headroom/providers/litellm.py +++ b/headroom/providers/litellm.py @@ -25,6 +25,7 @@ import logging import os from typing import Any +from headroom.pricing.litellm_pricing import estimate_cost_from_tokens from headroom.tokenizers import EstimatingTokenCounter from .base import Provider, TokenCounter @@ -240,19 +241,13 @@ class LiteLLMProvider(Provider): Returns: Estimated cost in USD, or None if pricing unknown. """ - try: - # LiteLLM's cost calculation - cost = litellm.completion_cost( - model=model, - prompt="", # We're using token counts directly - completion="", - prompt_tokens=input_tokens, - completion_tokens=output_tokens, - ) - return cost - except Exception as e: - logger.debug(f"LiteLLM cost estimation failed for {model}: {e}") - return None + # LiteLLM's cost calculation, from token counts directly. + return estimate_cost_from_tokens( + model, + input_tokens=input_tokens, + output_tokens=output_tokens, + cached_tokens=cached_tokens, + ) @classmethod def list_supported_providers(cls) -> list[str]: diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index 891b13559..bb61e9300 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -16,6 +16,7 @@ from functools import lru_cache from typing import Any, cast from headroom import paths as _paths +from headroom.pricing.litellm_pricing import estimate_cost_from_tokens from headroom.tokenizers.base import coerce_countable_text, count_content_blocks from .base import Provider, TokenCounter @@ -637,20 +638,16 @@ class OpenAIProvider(Provider): Returns: Estimated cost in USD, or None if pricing unknown. """ - # Try LiteLLM first (most up-to-date pricing) - litellm = _get_litellm_module() - if litellm is not None: - try: - # LiteLLM uses per-token pricing, returns total cost - cost = litellm.completion_cost( - model=model, - prompt_tokens=input_tokens, - completion_tokens=output_tokens, - ) - if cost is not None and cost > 0: - return float(cost) - except Exception: - pass # Fall through to manual pricing + # Try LiteLLM first (most up-to-date pricing, and it knows each model's + # real cached-input rate rather than the manual path's flat estimate) + cost = estimate_cost_from_tokens( + model, + input_tokens=input_tokens, + output_tokens=output_tokens, + cached_tokens=cached_tokens, + ) + if cost is not None and cost > 0: + return float(cost) # Fall back to hardcoded pricing return self._estimate_cost_manual(input_tokens, output_tokens, model, cached_tokens) diff --git a/tests/test_providers/test_anthropic.py b/tests/test_providers/test_anthropic.py index 833a6ffc9..d4a1b9ef7 100644 --- a/tests/test_providers/test_anthropic.py +++ b/tests/test_providers/test_anthropic.py @@ -35,6 +35,154 @@ class TestAnthropicModelSanitization: } +class TestContext1MSuffix: + """`[1m]` is a 1M-context tier request, not just an ANSI artifact (#1158). + + Claude Code appends `[1m]` to a model id and only then sends the + `context-1m` beta header, so the real upstream window is 1M even when the + base model defaults to 200K. The suffix must still be stripped off the wire + (upstream rejects it, #2027) but must not be lost before we size the budget. + """ + + @pytest.fixture + def provider(self): + from headroom.providers.anthropic import AnthropicProvider + + return AnthropicProvider() + + def test_1m_suffix_is_detected(self): + from headroom.providers.anthropic import has_context_1m_suffix + + assert has_context_1m_suffix("claude-sonnet-4-5[1m]") + assert has_context_1m_suffix("claude-sonnet-4-5[1m][1m]") + assert not has_context_1m_suffix("claude-sonnet-4-5") + + def test_ansi_artifacts_are_not_mistaken_for_a_tier_request(self): + from headroom.providers.anthropic import has_context_1m_suffix + + # A dangling reset, a compound style, and a real escape sequence are + # terminal noise -- none of them means "give me 1M". + assert not has_context_1m_suffix("claude-sonnet-4-5[0m]") + assert not has_context_1m_suffix("claude-sonnet-4-5[1;32m]") + assert not has_context_1m_suffix("\x1b[1mclaude-sonnet-4-5\x1b[0m") + + def test_1m_suffix_raises_a_200k_model_to_1m(self, provider): + # The regression: sanitizing before the lookup resolved this to the + # base model's 200K window, so a 1M request was budgeted at 1/5 size. + assert provider.get_context_limit("claude-sonnet-4-5") == 200_000 + assert provider.get_context_limit("claude-sonnet-4-5[1m]") == 1_000_000 + + def test_1m_suffix_never_lowers_an_already_larger_window(self, provider): + # max(), not a flat assignment: a base model wider than 1M keeps its own. + assert provider.get_context_limit("claude-opus-5[1m]") >= 1_000_000 + + def test_ansi_artifact_does_not_inflate_the_window(self, provider): + assert provider.get_context_limit("claude-sonnet-4-5[0m]") == 200_000 + assert provider.get_context_limit("\x1b[1mclaude-sonnet-4-5\x1b[0m") == 200_000 + + def test_wire_model_id_still_drops_the_suffix(self): + # Upstream rejects `[1m]`; the tier fix must not regress #2027. + from headroom.providers.anthropic import sanitize_anthropic_model_id + + assert sanitize_anthropic_model_id("claude-sonnet-4-5[1m]") == "claude-sonnet-4-5" + + +class TestLongContextPricing: + """Anthropic's long-context premium above a 200K prompt. + + On the Sonnet 4 / 4.5 family a prompt over 200K re-prices the *whole* + request -- input, output and cache alike -- at input 2x, output 1.5x, + cache 2x. Both the LiteLLM path and the manual fallback must apply it, or + Headroom under-reports the cost of exactly the sessions `[1m]` unlocks. + """ + + @pytest.fixture + def provider(self): + from headroom.providers.anthropic import AnthropicProvider + + return AnthropicProvider() + + @pytest.fixture + def manual_provider(self, monkeypatch): + """Provider with the LiteLLM path disabled, exercising the fallback.""" + import headroom.providers.anthropic as anthropic_module + + monkeypatch.setattr(anthropic_module, "estimate_cost_from_tokens", lambda *a, **k: None) + return anthropic_module.AnthropicProvider() + + # 100K in / 5K out -> 100K*$3 + 5K*$15 = $0.375 + # 300K in / 5K out -> 300K*$6 + 5K*$22.5 = $1.9125 (premium) + # 300K in of which 150K cached, 5K out + # -> 150K*$6 + 150K*$0.60 + 5K*$22.5 = $1.1025 + _CASES = [ + (100_000, 5_000, 0, 0.3750), + (300_000, 5_000, 0, 1.9125), + (300_000, 5_000, 150_000, 1.1025), + ] + + @pytest.mark.parametrize(("input_tokens", "output_tokens", "cached_tokens", "expected"), _CASES) + def test_litellm_path(self, provider, input_tokens, output_tokens, cached_tokens, expected): + cost = provider.estimate_cost( + input_tokens, output_tokens, "claude-sonnet-4-5", cached_tokens + ) + assert cost == pytest.approx(expected, rel=1e-4) + + @pytest.mark.parametrize(("input_tokens", "output_tokens", "cached_tokens", "expected"), _CASES) + def test_manual_fallback_matches_litellm( + self, manual_provider, input_tokens, output_tokens, cached_tokens, expected + ): + cost = manual_provider.estimate_cost( + input_tokens, output_tokens, "claude-sonnet-4-5", cached_tokens + ) + assert cost == pytest.approx(expected, rel=1e-4) + + def test_untiered_model_is_not_charged_a_premium(self, manual_provider): + # Opus is flat-rated across its whole window: 300K*$5 + 5K*$25 = $1.625. + cost = manual_provider.estimate_cost(300_000, 5_000, "claude-opus-4-5-20251101", 0) + assert cost == pytest.approx(1.625, rel=1e-4) + + def test_premium_applies_only_above_the_threshold(self, manual_provider): + at = manual_provider.estimate_cost(200_000, 0, "claude-sonnet-4-5", 0) + just_over = manual_provider.estimate_cost(200_001, 0, "claude-sonnet-4-5", 0) + assert at == pytest.approx(0.60, rel=1e-4) # 200K * $3 + assert just_over == pytest.approx(1.2000, rel=1e-3) # re-priced at $6 + + def test_1m_suffix_request_is_priced_at_the_premium(self, manual_provider): + # The two halves of this PR meeting: `[1m]` unlocks the window, and a + # session that fills it is billed at the long-context rate. + assert manual_provider.get_context_limit("claude-sonnet-4-5[1m]") == 1_000_000 + cost = manual_provider.estimate_cost(300_000, 5_000, "claude-sonnet-4-5[1m]", 0) + assert cost == pytest.approx(1.9125, rel=1e-4) + + +class TestLiteLLMCostHelper: + """The shared helper each provider now uses for LiteLLM-backed pricing. + + It replaces a `litellm.completion_cost(prompt_tokens=...)` call that had + stopped accepting those kwargs and raised TypeError on every invocation. + """ + + def test_returns_none_for_unknown_model(self): + from headroom.pricing.litellm_pricing import estimate_cost_from_tokens + + assert estimate_cost_from_tokens("no-such-model-xyz", 1000, 1000) is None + + def test_prices_a_known_model(self): + from headroom.pricing.litellm_pricing import estimate_cost_from_tokens + + # gpt-4o: $2.50/1M in, $10/1M out -> 100K in + 5K out = $0.30 + assert estimate_cost_from_tokens("gpt-4o", 100_000, 5_000) == pytest.approx(0.30, rel=1e-4) + + def test_input_tokens_are_cache_inclusive(self): + from headroom.pricing.litellm_pricing import estimate_cost_from_tokens + + # The cached portion is a subset of input_tokens, not additional to it, + # so a fully-cached prompt costs strictly less than an uncached one. + uncached = estimate_cost_from_tokens("gpt-4o", 100_000, 5_000) + cached = estimate_cost_from_tokens("gpt-4o", 100_000, 5_000, cached_tokens=50_000) + assert cached < uncached + + class TestAnthropicTokenCounting: @pytest.fixture def anthropic_provider(self): @@ -116,13 +264,16 @@ class TestAnthropicCostEstimation: return AnthropicProvider() def test_estimate_cost_basic(self, anthropic_provider): + # Probed at 100K, below the 200K long-context threshold: a 1M-token + # probe would cross it and bill at the premium rate, which is a + # separate property (covered by TestLongContextPricing). cost = anthropic_provider.estimate_cost( - input_tokens=1000000, + input_tokens=100_000, output_tokens=0, model="claude-3-5-sonnet-20241022", ) # $3.00 per 1M input - assert cost == pytest.approx(3.00, rel=0.1) + assert cost == pytest.approx(0.30, rel=0.1) def test_pricing_lookup_strips_ansi_model_suffix(self, anthropic_provider): assert anthropic_provider._get_pricing("claude-opus-4-7[1m]") == ( diff --git a/tests/test_providers/test_universal.py b/tests/test_providers/test_universal.py index b840837fc..aa3f337d4 100644 --- a/tests/test_providers/test_universal.py +++ b/tests/test_providers/test_universal.py @@ -443,22 +443,13 @@ class TestLiteLLMProvider: "output-model": {"max_output_tokens": 6000}, }[model], ) + # Cost now resolves through the shared pricing helper rather than a + # direct `litellm.completion_cost` call, so patch that seam. The helper + # returns None (not an exception) for a model LiteLLM can't price. monkeypatch.setattr( litellm_module, - "litellm", - type( - "LiteLLM", - (), - { - "completion_cost": staticmethod( - lambda **kwargs: ( - 1.23 - if kwargs["model"] == "priced-model" - else (_ for _ in ()).throw(RuntimeError("missing price")) - ) - ) - }, - )(), + "estimate_cost_from_tokens", + lambda model, **kwargs: 1.23 if model == "priced-model" else None, ) provider = litellm_module.LiteLLMProvider()