diff --git a/headroom/pricing/litellm_model_resolution.py b/headroom/pricing/litellm_model_resolution.py index c006f64d9..0a9242001 100644 --- a/headroom/pricing/litellm_model_resolution.py +++ b/headroom/pricing/litellm_model_resolution.py @@ -2,9 +2,15 @@ from __future__ import annotations +import re from collections.abc import Callable from dataclasses import dataclass +# Vertex AI appends @YYYYMMDD version tags to model names at runtime +# (e.g. "claude-haiku-4-5@20251001"). LiteLLM's database stores bare +# names without version suffixes, so we strip the suffix before lookup. +_VERTEX_VERSION_SUFFIX_RE = re.compile(r"@\d{8}$") + @dataclass(frozen=True, slots=True) class LiteLLMModelPrefixRule: @@ -55,16 +61,40 @@ PRICE_LOOKUP_PROVIDER_PREFIXES: tuple[str, ...] = ( ) +def _strip_vertex_version_suffix(model: str) -> str: + """Strip Vertex @YYYYMMDD version suffix if present.""" + return _VERTEX_VERSION_SUFFIX_RE.sub("", model) + + def resolution_candidates(model: str) -> tuple[str, ...]: """Return ordered LiteLLM keys to try for cost-per-token resolution.""" candidates = [model] - candidates.extend( - candidate - for rule in MODEL_PREFIX_RULES - for candidate in (rule.candidate_for(model),) - if candidate is not None - ) - alias = MODEL_ALIASES.get(model) + + # If the model has a Vertex @YYYYMMDD version suffix, also try the bare + # name. Vertex appends these at runtime; LiteLLM stores bare names only. + bare = _strip_vertex_version_suffix(model) + is_vertex_versioned = bare != model + if is_vertex_versioned: + candidates.append(bare) + + # Apply prefix rules to both the original and bare name so that e.g. + # "anthropic/claude-haiku-4-5" is tried after "claude-haiku-4-5". + for m in dict.fromkeys([model, bare]): + candidates.extend( + candidate + for rule in MODEL_PREFIX_RULES + for candidate in (rule.candidate_for(m),) + if candidate is not None + ) + + # Only add vertex_ai/ candidates for models with @YYYYMMDD suffix — + # these are known Vertex-routed models. Non-versioned models should not + # get vertex_ai/ candidates to avoid matching wrong pricing tier. + if is_vertex_versioned: + for m in dict.fromkeys([model, bare]): + candidates.append(f"vertex_ai/{m}") + + alias = MODEL_ALIASES.get(model) or MODEL_ALIASES.get(bare) if alias: candidates.append(alias) return tuple(dict.fromkeys(candidates)) @@ -86,11 +116,27 @@ def unwrapped_model_forms(model: str) -> tuple[str, ...]: def pricing_lookup_candidates(model: str) -> tuple[str, ...]: """Return ordered LiteLLM model_cost keys to try for pricing lookup.""" + bare = _strip_vertex_version_suffix(model) + is_vertex_versioned = bare != model + candidates = [model] - candidates.extend(f"{prefix}{model}" for prefix in PRICE_LOOKUP_PROVIDER_PREFIXES) + if is_vertex_versioned: + candidates.append(bare) + + # Try all provider prefixes for both the original and bare name. + for m in dict.fromkeys([model, bare]): + candidates.extend(f"{prefix}{m}" for prefix in PRICE_LOOKUP_PROVIDER_PREFIXES) + # Unwrapped forms come after the prefixed ones so existing precedence is # unchanged for names that already resolved. - candidates.extend(unwrapped_model_forms(model)) + for m in dict.fromkeys([model, bare]): + candidates.extend(unwrapped_model_forms(m)) + + # Only add vertex_ai/ candidates for models with @YYYYMMDD suffix. + if is_vertex_versioned: + for m in dict.fromkeys([model, bare]): + candidates.append(f"vertex_ai/{m}") + candidates.extend( alias for candidate in tuple(candidates) if (alias := MODEL_ALIASES.get(candidate)) ) diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index cff6be4b2..b42ac18b0 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -213,7 +213,7 @@ def build_prefix_cache_stats( # Match model to provider _openai_prefixes = ("gpt", "o1", "o3", "o4") is_match = ( - (provider == "anthropic" and "claude" in model_name) + (provider in ("anthropic", "vertex: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) diff --git a/tests/test_pricing_litellm_model_resolution.py b/tests/test_pricing_litellm_model_resolution.py index c2098a89b..dcfb21485 100644 --- a/tests/test_pricing_litellm_model_resolution.py +++ b/tests/test_pricing_litellm_model_resolution.py @@ -3,6 +3,7 @@ from __future__ import annotations from headroom.pricing.litellm_model_resolution import ( MODEL_ALIASES, LiteLLMModelPrefixRule, + _strip_vertex_version_suffix, pricing_lookup_candidates, resolution_candidates, resolve_litellm_model_name, @@ -21,11 +22,11 @@ def test_resolution_candidates_try_bare_then_matching_prefix_then_alias() -> Non assert resolution_candidates("MiniMax-M3") == ("MiniMax-M3", "minimax/MiniMax-M3") retired = "claude-3-5-sonnet-20241022" - assert resolution_candidates(retired) == ( - retired, - f"anthropic/{retired}", - MODEL_ALIASES[retired], - ) + candidates = resolution_candidates(retired) + assert candidates[0] == retired + assert f"anthropic/{retired}" in candidates + assert f"vertex_ai/{retired}" not in candidates # no @YYYYMMDD suffix = not Vertex + assert MODEL_ALIASES[retired] in candidates def test_pricing_lookup_candidates_include_provider_prefixes_and_aliases() -> None: @@ -59,3 +60,64 @@ def test_resolve_litellm_model_name_returns_first_known_candidate() -> None: def test_resolve_litellm_model_name_returns_original_when_unknown() -> None: assert resolve_litellm_model_name("mystery-model", lambda _: False) == "mystery-model" + + +def test_strip_vertex_version_suffix() -> None: + assert _strip_vertex_version_suffix("claude-haiku-4-5@20251001") == "claude-haiku-4-5" + assert _strip_vertex_version_suffix("claude-opus-4@20250514") == "claude-opus-4" + assert _strip_vertex_version_suffix("claude-sonnet-4-6") == "claude-sonnet-4-6" + assert _strip_vertex_version_suffix("claude-sonnet-4-20250514") == "claude-sonnet-4-20250514" + + +def test_resolution_candidates_vertex_versioned_models() -> None: + # Vertex appends @YYYYMMDD — bare name and vertex_ai/ must be candidates + candidates = resolution_candidates("claude-haiku-4-5@20251001") + assert "claude-haiku-4-5" in candidates + assert "anthropic/claude-haiku-4-5" in candidates + assert "vertex_ai/claude-haiku-4-5" in candidates # vertex_ai/ only for versioned + + candidates = resolution_candidates("claude-opus-4@20250514") + assert "claude-opus-4" in candidates + assert "anthropic/claude-opus-4" in candidates + assert "vertex_ai/claude-opus-4" in candidates + + # Non-versioned names should NOT get vertex_ai/ candidates + candidates = resolution_candidates("claude-sonnet-4-6") + assert candidates[0] == "claude-sonnet-4-6" + assert "vertex_ai/claude-sonnet-4-6" not in candidates + assert "anthropic/claude-sonnet-4-6" in candidates + + +def test_pricing_lookup_candidates_vertex_versioned_models() -> None: + candidates = pricing_lookup_candidates("claude-haiku-4-5@20251001") + # Bare name and vertex_ai/ prefix must both be candidates + assert "claude-haiku-4-5" in candidates + assert "vertex_ai/claude-haiku-4-5" in candidates + assert "anthropic/claude-haiku-4-5" in candidates + + candidates = pricing_lookup_candidates("claude-opus-4@20250514") + assert "claude-opus-4" in candidates + assert "vertex_ai/claude-opus-4" in candidates + + # Non-versioned names should NOT get vertex_ai/ pricing candidates + candidates = pricing_lookup_candidates("claude-sonnet-4-6") + assert "vertex_ai/claude-sonnet-4-6" not in candidates + assert "anthropic/claude-sonnet-4-6" in candidates + + +def test_vertex_versioned_model_resolves_to_known_key() -> None: + # Simulate LiteLLM knowing the bare model name (not the versioned one) + known = {"claude-haiku-4-5", "anthropic/claude-sonnet-4-6"} + assert ( + resolve_litellm_model_name("claude-haiku-4-5@20251001", known.__contains__) + == "claude-haiku-4-5" + ) + assert ( + resolve_litellm_model_name("claude-sonnet-4-6", known.__contains__) + == "anthropic/claude-sonnet-4-6" + ) + # Unknown versioned model falls back to original + assert ( + resolve_litellm_model_name("claude-unknown@20251001", lambda _: False) + == "claude-unknown@20251001" + )