diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 6860fed58..1d597b544 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -226,6 +226,39 @@ def _fetch_bedrock_inference_profiles( return model_map +def _parse_bedrock_model_overrides(raw: str | None) -> dict[str, str]: + """Parse the ``HEADROOM_BEDROCK_MODEL_MAP`` operator override. + + AWS discovery keys the model map by the *normalized model name*, so it + cannot disambiguate application inference profiles that share one + underlying model — e.g. a team where ``claude-sonnet-5-kenneth`` and + ``claude-sonnet-5-jeremy`` both resolve to ``claude-sonnet-5``. When you + need requests billed to a *specific* application profile (per-user cost + attribution), pin the mapping explicitly here. The plain name Claude Code + sends (kept plain so tool-search deferral stays on) resolves to your ARN. + + Format: comma-separated ``name=target`` pairs, where ``target`` is an + application-inference-profile ARN (routed via the converse endpoint) or + any LiteLLM model string. Whitespace around pairs is ignored; blank + entries are skipped. + + HEADROOM_BEDROCK_MODEL_MAP="claude-sonnet-5=arn:aws:bedrock:...:application-inference-profile/x57j1esjrt66,claude-opus-4-8=arn:aws:bedrock:...:application-inference-profile/3dy9ytxuq2ci" + """ + overrides: dict[str, str] = {} + if not raw: + return overrides + for pair in raw.split(","): + pair = pair.strip() + if not pair or "=" not in pair: + continue + name, _, target = pair.partition("=") + name = name.strip() + target = target.strip() + if name and target: + overrides[name] = target + return overrides + + def _normalize_bedrock_profile_id(profile_id: str) -> str | None: """Extract standard Anthropic model name from Bedrock profile ID. @@ -248,8 +281,10 @@ def _normalize_bedrock_profile_id(profile_id: str) -> str | None: if profile_id.startswith("bedrock/"): profile_id = profile_id[8:] - # Strip region prefix (us., eu., apac., au.) - for prefix in ["us.", "eu.", "apac.", "au."]: + # Strip region prefix (us., eu., apac., au.) or the newer "global." + # cross-region prefix used by current-gen profiles (e.g. + # "global.anthropic.claude-sonnet-4-6"). + for prefix in ["us.", "eu.", "apac.", "au.", "global."]: if profile_id.startswith(prefix): profile_id = profile_id[len(prefix) :] break @@ -262,8 +297,12 @@ def _normalize_bedrock_profile_id(profile_id: str) -> str | None: if not profile_id.startswith("claude"): return None - # Strip version suffix (-v1:0, -v2:0, etc.) - normalized = re.sub(r"-v\d+:\d+$", "", profile_id) + # Strip version suffix. Legacy dated profiles use "-v1:0" / "-v2:0"; + # newer undated profiles use a bare "-v1" (no colon/revision) or carry + # no version suffix at all (e.g. "claude-opus-4-8"). Match all three + # shapes so undated current-gen profiles normalize instead of + # silently falling out of the resolvable model map. + normalized = re.sub(r"-v\d+(?::\d+)?$", "", profile_id) return normalized if normalized else None @@ -491,6 +530,20 @@ class LiteLLMBackend(Backend): else: self._model_map = self._config.model_map + # Operator override map (all providers; only meaningful for Bedrock + # today). Lets you pin a plain model name to a specific target the + # AWS discovery can't disambiguate — e.g. a per-user application + # inference profile ARN for cost attribution. See + # `_parse_bedrock_model_overrides`. + self._model_overrides = _parse_bedrock_model_overrides( + os.environ.get("HEADROOM_BEDROCK_MODEL_MAP") + ) + if self._model_overrides: + logger.info( + f"Loaded {len(self._model_overrides)} Bedrock model override(s) " + f"from HEADROOM_BEDROCK_MODEL_MAP: {sorted(self._model_overrides)}" + ) + logger.info(f"LiteLLM backend initialized (provider={provider}, region={region})") @property @@ -507,6 +560,19 @@ class LiteLLMBackend(Backend): - "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" (LiteLLM format) - "arn:aws:bedrock:...:application-inference-profile/..." (application inference profile) """ + # Operator override wins over everything — an explicit pin the AWS + # discovery cannot express (e.g. a per-user application inference + # profile). Keyed by the plain name Claude Code sends. + override = self._model_overrides.get(anthropic_model) + if override: + if override.startswith("arn:aws:"): + # Application inference profile ARNs must use the converse + # route — the invoke route rejects ARNs with HTTP 400. + return f"bedrock/converse/{override}" + if override.startswith(f"{self.provider}/"): + return override + return f"{self.provider}/{override}" + # Check direct mapping first if anthropic_model in self._model_map: return self._model_map[anthropic_model] diff --git a/tests/test_bedrock_region.py b/tests/test_bedrock_region.py index 46df34b2d..f96cff1a6 100644 --- a/tests/test_bedrock_region.py +++ b/tests/test_bedrock_region.py @@ -347,6 +347,33 @@ class TestBedrockModelMapping: result = backend.map_model_id("claude-sonnet-4-5-20250929") assert result == "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0" + def test_override_pins_plain_name_to_app_profile_arn(self, monkeypatch): + """HEADROOM_BEDROCK_MODEL_MAP pins a plain name to a specific app + profile ARN and routes it via the converse endpoint, winning over + the discovered map.""" + arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66" + monkeypatch.setenv("HEADROOM_BEDROCK_MODEL_MAP", f"claude-sonnet-5={arn}") + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + # Discovery also has a system-defined sonnet-5; the override must win. + return_value={"claude-sonnet-5": "bedrock/global.anthropic.claude-sonnet-5"}, + ): + backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1") + assert backend.map_model_id("claude-sonnet-5") == f"bedrock/converse/{arn}" + + def test_override_absent_falls_through_to_discovery(self, monkeypatch): + """A model not pinned in the override map resolves via discovery.""" + arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66" + monkeypatch.setenv("HEADROOM_BEDROCK_MODEL_MAP", f"claude-sonnet-5={arn}") + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={"claude-opus-4-8": "bedrock/global.anthropic.claude-opus-4-8"}, + ): + backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1") + assert backend.map_model_id("claude-opus-4-8") == ( + "bedrock/global.anthropic.claude-opus-4-8" + ) + # ============================================================================= # Normalize Bedrock Profile ID (edge cases) @@ -390,6 +417,71 @@ class TestNormalizeBedrockProfileId: "claude-sonnet-4-20250514" ) + def test_global_prefix_no_version_suffix(self): + # Current-gen cross-region profile: "global." prefix, no version + # suffix at all. Must strip the prefix and keep the bare name. + assert _normalize_bedrock_profile_id("global.anthropic.claude-opus-4-8") == ( + "claude-opus-4-8" + ) + + def test_global_prefix_bare_v_suffix(self): + # "global." prefix with an undated "-v1" (no ":revision") suffix. + assert _normalize_bedrock_profile_id("global.anthropic.claude-opus-4-6-v1") == ( + "claude-opus-4-6" + ) + + def test_global_prefix_dated_full_suffix(self): + # "global." prefix with the legacy dated "-vN:M" suffix. + assert ( + _normalize_bedrock_profile_id("global.anthropic.claude-haiku-4-5-20251001-v1:0") + == "claude-haiku-4-5-20251001" + ) + + def test_global_prefix_next_gen_names(self): + assert ( + _normalize_bedrock_profile_id("global.anthropic.claude-sonnet-5") == "claude-sonnet-5" + ) + assert _normalize_bedrock_profile_id("global.anthropic.claude-fable-5") == "claude-fable-5" + + +# ============================================================================= +# HEADROOM_BEDROCK_MODEL_MAP operator override parsing +# ============================================================================= + + +class TestParseBedrockModelOverrides: + """Test the HEADROOM_BEDROCK_MODEL_MAP override parser.""" + + def test_none_and_empty_yield_empty(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + assert _parse_bedrock_model_overrides(None) == {} + assert _parse_bedrock_model_overrides("") == {} + assert _parse_bedrock_model_overrides(" ") == {} + + def test_single_pair(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66" + assert _parse_bedrock_model_overrides(f"claude-sonnet-5={arn}") == {"claude-sonnet-5": arn} + + def test_multiple_pairs_and_whitespace(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + raw = " claude-sonnet-5=arn:a , claude-opus-4-8=arn:b " + assert _parse_bedrock_model_overrides(raw) == { + "claude-sonnet-5": "arn:a", + "claude-opus-4-8": "arn:b", + } + + def test_skips_malformed_entries(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + # Missing "=" and blank segments are skipped, valid pairs survive. + assert _parse_bedrock_model_overrides("garbage,,claude-sonnet-5=arn:a,=noname") == { + "claude-sonnet-5": "arn:a", + } + # ============================================================================= # Named profile forwarded to acompletion kwargs