From b0ece04b9b4dba714a9034f13c1842448c1af110 Mon Sep 17 00:00:00 2001 From: Yitong Li Date: Tue, 24 Mar 2026 00:13:00 +0800 Subject: [PATCH] fix(bedrock): add EU/AP region support with graceful fallback Three issues fixed: 1. _fetch_bedrock_inference_profiles crashed the proxy on startup when boto3 was missing or the AWS API call failed (wrong credentials, permissions, network). Now catches exceptions and falls back to a static model map. 2. map_model_id produced invalid Bedrock model IDs for unmapped models. Bare names like 'claude-sonnet-4-20250514' became 'bedrock/claude-sonnet-4-20250514' which is not a valid Bedrock identifier. Now constructs region-prefixed IDs like 'bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0'. 3. No static fallback existed (_BEDROCK_MODEL_MAP was empty). Added _build_bedrock_fallback_map() that generates region-aware model IDs for all GA Claude models (us./eu./apac. prefixes). Closes #28 Tests: 27 new tests covering region prefix mapping, static fallback map, graceful degradation, and model ID mapping for EU/AP/US regions. --- headroom/backends/litellm.py | 148 ++++++++++++--- tests/test_bedrock_region.py | 352 +++++++++++++++++++++++++++++++++++ 2 files changed, 472 insertions(+), 28 deletions(-) create mode 100644 tests/test_bedrock_region.py diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 9c9deed15..ddcf0b5f8 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -55,6 +55,75 @@ class ProviderConfig: # Cache for dynamically fetched inference profiles _bedrock_profiles_cache: dict[str, dict[str, str]] = {} # region -> model_map +# Region prefix used in cross-region Bedrock inference profile IDs. +# EU regions use "eu.", AP regions use "apac.", US (and everything else) use "us.". +_BEDROCK_REGION_PREFIXES: dict[str, str] = { + "eu": "eu", + "ap": "apac", +} + + +def _bedrock_region_prefix(region: str) -> str: + """Return the inference-profile region prefix for an AWS region. + + AWS Bedrock cross-region inference profiles are prefixed with a + geographic tag: ``us.``, ``eu.``, or ``apac.``. This helper maps + an AWS region name (e.g. ``eu-west-1``) to the correct prefix. + + >>> _bedrock_region_prefix("us-east-1") + 'us' + >>> _bedrock_region_prefix("eu-central-1") + 'eu' + >>> _bedrock_region_prefix("ap-southeast-1") + 'apac' + """ + for key, prefix in _BEDROCK_REGION_PREFIXES.items(): + if region.startswith(key): + return prefix + return "us" + + +def _build_bedrock_fallback_map(region: str) -> dict[str, str]: + """Build a static Bedrock model map using the region prefix. + + When ``_fetch_bedrock_inference_profiles`` cannot reach the AWS API + (wrong credentials, network error, permissions, etc.) we fall back + to this map so that the proxy can still route requests. The map + covers all currently GA Claude models on Bedrock. + """ + prefix = _bedrock_region_prefix(region) + + # Base model IDs without region prefix + _CLAUDE_MODELS = [ + # Claude 4.6 + ("claude-opus-4-6", "anthropic.claude-opus-4-6-v1:0"), + ("claude-sonnet-4-6", "anthropic.claude-sonnet-4-6-v1:0"), + # Claude 4.5 + ("claude-sonnet-4-5-20250929", "anthropic.claude-sonnet-4-5-20250929-v1:0"), + ("claude-opus-4-5-20251101", "anthropic.claude-opus-4-5-20251101-v1:0"), + # Claude 4.1 + ("claude-opus-4-1-20250805", "anthropic.claude-opus-4-1-20250805-v1:0"), + # Claude 4 + ("claude-sonnet-4-20250514", "anthropic.claude-sonnet-4-20250514-v1:0"), + ("claude-opus-4-20250514", "anthropic.claude-opus-4-20250514-v1:0"), + # Claude 3.7 + ("claude-3-7-sonnet-20250219", "anthropic.claude-3-7-sonnet-20250219-v1:0"), + # Claude 3.5 + ("claude-3-5-sonnet-20241022", "anthropic.claude-3-5-sonnet-20241022-v2:0"), + ("claude-3-5-sonnet-20240620", "anthropic.claude-3-5-sonnet-20240620-v1:0"), + ("claude-3-5-haiku-20241022", "anthropic.claude-3-5-haiku-20241022-v1:0"), + # Claude 3 + ("claude-3-opus-20240229", "anthropic.claude-3-opus-20240229-v1:0"), + ("claude-3-sonnet-20240229", "anthropic.claude-3-sonnet-20240229-v1:0"), + ("claude-3-haiku-20240307", "anthropic.claude-3-haiku-20240307-v1:0"), + # Haiku 4.5 + ("claude-haiku-4-5-20251001", "anthropic.claude-haiku-4-5-20251001-v1:0"), + ] + + return { + name: f"bedrock/{prefix}.{model_id}" for name, model_id in _CLAUDE_MODELS + } + def _fetch_bedrock_inference_profiles(region: str | None) -> dict[str, str]: """Fetch available Bedrock inference profiles from AWS API. @@ -62,19 +131,16 @@ def _fetch_bedrock_inference_profiles(region: str | None) -> dict[str, str]: Uses boto3 list_inference_profiles() to get all available profiles for the given region, then builds a model map. + If the API call fails (wrong credentials, network error, permission + denied, etc.) the function logs a warning and returns a static + fallback map so the proxy can still start. + Args: region: AWS region (e.g., "us-east-1", "eu-central-1") Returns: Model map: anthropic_model_name -> bedrock inference profile ID - - Raises: - ImportError: If boto3 is not installed - Exception: If API call fails """ - - import boto3 - region = region or "us-east-1" # Check cache first @@ -83,36 +149,54 @@ def _fetch_bedrock_inference_profiles(region: str | None) -> dict[str, str]: model_map: dict[str, str] = {} - bedrock_client = boto3.client("bedrock", region_name=region) - response = bedrock_client.list_inference_profiles(typeEquals="SYSTEM_DEFINED") - - for profile in response.get("inferenceProfileSummaries", []): - profile_id = profile.get("inferenceProfileId", "") - - # Only process Anthropic Claude profiles - if "anthropic" not in profile_id.lower(): - continue - - # Extract the standard model name from the profile ID - # e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0" -> "claude-sonnet-4-20250514" - normalized = _normalize_bedrock_profile_id(profile_id) - if normalized: - model_map[normalized] = f"bedrock/{profile_id}" - - # Handle pagination if needed - while response.get("nextToken"): - response = bedrock_client.list_inference_profiles( - typeEquals="SYSTEM_DEFINED", nextToken=response["nextToken"] + try: + import boto3 + except ImportError: + logger.warning( + "boto3 is not installed — using static Bedrock model map. " + "Install boto3 for dynamic model discovery: pip install boto3" ) + model_map = _build_bedrock_fallback_map(region) + _bedrock_profiles_cache[region] = model_map + return model_map + + try: + bedrock_client = boto3.client("bedrock", region_name=region) + response = bedrock_client.list_inference_profiles(typeEquals="SYSTEM_DEFINED") + for profile in response.get("inferenceProfileSummaries", []): profile_id = profile.get("inferenceProfileId", "") + + # Only process Anthropic Claude profiles if "anthropic" not in profile_id.lower(): continue + + # Extract the standard model name from the profile ID + # e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0" -> "claude-sonnet-4-20250514" normalized = _normalize_bedrock_profile_id(profile_id) if normalized: model_map[normalized] = f"bedrock/{profile_id}" - logger.info(f"Fetched {len(model_map)} Bedrock inference profiles for region {region}") + # Handle pagination if needed + while response.get("nextToken"): + response = bedrock_client.list_inference_profiles( + typeEquals="SYSTEM_DEFINED", nextToken=response["nextToken"] + ) + for profile in response.get("inferenceProfileSummaries", []): + profile_id = profile.get("inferenceProfileId", "") + if "anthropic" not in profile_id.lower(): + continue + normalized = _normalize_bedrock_profile_id(profile_id) + if normalized: + model_map[normalized] = f"bedrock/{profile_id}" + + logger.info(f"Fetched {len(model_map)} Bedrock inference profiles for region {region}") + except Exception as e: + logger.warning( + f"Failed to fetch Bedrock inference profiles for region {region}: {e}. " + "Using static fallback model map." + ) + model_map = _build_bedrock_fallback_map(region) # Cache the result _bedrock_profiles_cache[region] = model_map @@ -359,6 +443,14 @@ class LiteLLMBackend(Backend): if normalized and normalized in self._model_map: return self._model_map[normalized] + # Bedrock fallback: construct a valid region-prefixed model ID. + # Without this, bare model names like "claude-sonnet-4-20250514" + # would become "bedrock/claude-sonnet-4-20250514" which is not a + # valid Bedrock model identifier. + if "/" not in anthropic_model and anthropic_model.startswith("claude"): + region_prefix = _bedrock_region_prefix(self.region or "us-east-1") + return f"bedrock/{region_prefix}.anthropic.{anthropic_model}-v1:0" + # Pass-through providers: prepend provider prefix if self._config.pass_through: # If already has provider prefix, use as-is diff --git a/tests/test_bedrock_region.py b/tests/test_bedrock_region.py new file mode 100644 index 000000000..4b84c9e91 --- /dev/null +++ b/tests/test_bedrock_region.py @@ -0,0 +1,352 @@ +"""Tests for Bedrock region support and fallback model mapping. + +Ensures that EU, AP, and US regions all produce valid Bedrock model IDs, +and that the proxy degrades gracefully when boto3 is unavailable or the +AWS API call fails. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +pytest.importorskip("litellm") + +from headroom.backends.litellm import ( + LiteLLMBackend, + _bedrock_profiles_cache, + _bedrock_region_prefix, + _build_bedrock_fallback_map, + _fetch_bedrock_inference_profiles, + _normalize_bedrock_profile_id, +) + + +# ============================================================================= +# Region Prefix Mapping +# ============================================================================= + + +class TestBedrockRegionPrefix: + """Test AWS region -> inference profile prefix mapping.""" + + def test_us_regions(self): + assert _bedrock_region_prefix("us-east-1") == "us" + assert _bedrock_region_prefix("us-west-2") == "us" + + def test_eu_regions(self): + assert _bedrock_region_prefix("eu-central-1") == "eu" + assert _bedrock_region_prefix("eu-west-1") == "eu" + assert _bedrock_region_prefix("eu-west-3") == "eu" + assert _bedrock_region_prefix("eu-north-1") == "eu" + + def test_ap_regions(self): + assert _bedrock_region_prefix("ap-southeast-1") == "apac" + assert _bedrock_region_prefix("ap-northeast-1") == "apac" + + def test_unknown_region_defaults_to_us(self): + assert _bedrock_region_prefix("me-south-1") == "us" + assert _bedrock_region_prefix("sa-east-1") == "us" + + +# ============================================================================= +# Static Fallback Model Map +# ============================================================================= + + +class TestBuildBedrockFallbackMap: + """Test static fallback model map construction.""" + + def test_us_region_uses_us_prefix(self): + model_map = _build_bedrock_fallback_map("us-east-1") + assert "claude-sonnet-4-20250514" in model_map + assert model_map["claude-sonnet-4-20250514"] == ( + "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" + ) + + def test_eu_region_uses_eu_prefix(self): + model_map = _build_bedrock_fallback_map("eu-central-1") + assert "claude-sonnet-4-20250514" in model_map + assert model_map["claude-sonnet-4-20250514"] == ( + "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + ) + + def test_ap_region_uses_apac_prefix(self): + model_map = _build_bedrock_fallback_map("ap-southeast-1") + assert "claude-sonnet-4-20250514" in model_map + assert model_map["claude-sonnet-4-20250514"] == ( + "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0" + ) + + def test_all_models_present(self): + model_map = _build_bedrock_fallback_map("us-east-1") + expected_models = [ + "claude-opus-4-6", + "claude-sonnet-4-6", + "claude-sonnet-4-20250514", + "claude-opus-4-20250514", + "claude-3-7-sonnet-20250219", + "claude-3-5-sonnet-20241022", + "claude-3-5-haiku-20241022", + "claude-3-opus-20240229", + "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", + ] + for model in expected_models: + assert model in model_map, f"Missing model: {model}" + + def test_all_values_are_valid_bedrock_format(self): + """Every value must start with 'bedrock/' and contain 'anthropic.'.""" + for region in ("us-east-1", "eu-west-1", "ap-northeast-1"): + model_map = _build_bedrock_fallback_map(region) + for name, litellm_id in model_map.items(): + assert litellm_id.startswith("bedrock/"), ( + f"{name}: expected bedrock/ prefix, got {litellm_id}" + ) + assert "anthropic." in litellm_id, ( + f"{name}: expected anthropic. in id, got {litellm_id}" + ) + + +# ============================================================================= +# Fetch with Graceful Fallback +# ============================================================================= + + +class TestFetchBedrockInferenceProfiles: + """Test dynamic fetch with fallback on failure.""" + + def setup_method(self): + """Clear the cache before each test.""" + _bedrock_profiles_cache.clear() + + def test_fallback_when_boto3_import_fails(self): + """Should return static map when boto3 is not installed.""" + with patch.dict("sys.modules", {"boto3": None}): + # Force reimport failure + import importlib + + import headroom.backends.litellm as mod + + # Temporarily break boto3 import inside the function + original_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + def mock_import(name, *args, **kwargs): + if name == "boto3": + raise ImportError("No module named 'boto3'") + return original_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=mock_import): + _bedrock_profiles_cache.clear() + result = _fetch_bedrock_inference_profiles("eu-central-1") + + assert len(result) > 0 + # Should use EU prefix + assert result["claude-sonnet-4-20250514"] == ( + "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + ) + + def test_fallback_when_api_call_fails(self): + """Should return static map when list_inference_profiles raises.""" + mock_boto3 = MagicMock() + mock_client = MagicMock() + mock_client.list_inference_profiles.side_effect = Exception( + "AccessDeniedException: not authorized" + ) + mock_boto3.client.return_value = mock_client + + with patch("headroom.backends.litellm.boto3", mock_boto3, create=True): + # Patch the import inside the function + original_fn = _fetch_bedrock_inference_profiles.__code__ + _bedrock_profiles_cache.clear() + + # We need to actually test the function, so let's just use the + # mock_boto3 and make sure the function catches the exception + import builtins + real_import = builtins.__import__ + + def patched_import(name, *args, **kwargs): + if name == "boto3": + return mock_boto3 + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=patched_import): + _bedrock_profiles_cache.clear() + result = _fetch_bedrock_inference_profiles("eu-west-1") + + assert len(result) > 0 + # Should use EU prefix + for litellm_id in result.values(): + assert "eu.anthropic." in litellm_id + + def test_successful_fetch_uses_api_results(self): + """When API works, should use dynamic results (not fallback).""" + mock_boto3 = MagicMock() + mock_client = MagicMock() + mock_client.list_inference_profiles.return_value = { + "inferenceProfileSummaries": [ + {"inferenceProfileId": "eu.anthropic.claude-sonnet-4-20250514-v1:0"}, + {"inferenceProfileId": "eu.anthropic.claude-3-5-sonnet-20241022-v2:0"}, + {"inferenceProfileId": "eu.meta.llama-3-70b-v1:0"}, # non-Anthropic, should skip + ] + } + mock_boto3.client.return_value = mock_client + + import builtins + real_import = builtins.__import__ + + def patched_import(name, *args, **kwargs): + if name == "boto3": + return mock_boto3 + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=patched_import): + _bedrock_profiles_cache.clear() + result = _fetch_bedrock_inference_profiles("eu-central-1") + + assert len(result) == 2 + assert result["claude-sonnet-4-20250514"] == ( + "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + ) + assert result["claude-3-5-sonnet-20241022"] == ( + "bedrock/eu.anthropic.claude-3-5-sonnet-20241022-v2:0" + ) + + def test_caching_prevents_repeated_api_calls(self): + """Second call for same region should return cached result.""" + _bedrock_profiles_cache.clear() + _bedrock_profiles_cache["us-east-1"] = {"test": "bedrock/test-model"} + + result = _fetch_bedrock_inference_profiles("us-east-1") + assert result == {"test": "bedrock/test-model"} + + +# ============================================================================= +# LiteLLMBackend.map_model_id with EU Regions +# ============================================================================= + + +class TestBedrockModelMapping: + """Test model ID mapping for different regions.""" + + def setup_method(self): + _bedrock_profiles_cache.clear() + + def test_eu_region_maps_correctly(self): + """EU region should produce eu.anthropic.* model IDs.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={ + "claude-sonnet-4-20250514": "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", + }, + ): + backend = LiteLLMBackend(provider="bedrock", region="eu-central-1") + result = backend.map_model_id("claude-sonnet-4-20250514") + assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + + def test_us_region_maps_correctly(self): + """US region should produce us.anthropic.* model IDs.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={ + "claude-sonnet-4-20250514": "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0", + }, + ): + backend = LiteLLMBackend(provider="bedrock", region="us-west-2") + result = backend.map_model_id("claude-sonnet-4-20250514") + assert result == "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" + + def test_fallback_for_unknown_model_in_eu(self): + """Unknown models in EU should get eu.anthropic.* fallback, not bare 'bedrock/claude-...'.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={}, + ): + backend = LiteLLMBackend(provider="bedrock", region="eu-west-1") + result = backend.map_model_id("claude-sonnet-4-20250514") + assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + + def test_fallback_for_unknown_model_in_ap(self): + """Unknown models in AP should get apac.anthropic.* fallback.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={}, + ): + backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1") + result = backend.map_model_id("claude-3-5-haiku-20241022") + assert result == "bedrock/apac.anthropic.claude-3-5-haiku-20241022-v1:0" + + def test_bedrock_format_passthrough(self): + """Already-formatted Bedrock IDs should pass through unchanged.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={}, + ): + backend = LiteLLMBackend(provider="bedrock", region="eu-central-1") + model = "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + result = backend.map_model_id(model) + assert result == model + + def test_anthropic_dot_format_normalized(self): + """Raw Bedrock IDs like 'anthropic.claude-...-v1:0' should normalize and map.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={ + "claude-sonnet-4-20250514": "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", + }, + ): + backend = LiteLLMBackend(provider="bedrock", region="eu-central-1") + result = backend.map_model_id("anthropic.claude-sonnet-4-20250514-v1:0") + assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + + def test_region_prefixed_format_normalized(self): + """'eu.anthropic.claude-...-v1:0' should normalize and map.""" + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={ + "claude-sonnet-4-20250514": "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", + }, + ): + backend = LiteLLMBackend(provider="bedrock", region="eu-central-1") + result = backend.map_model_id("eu.anthropic.claude-sonnet-4-20250514-v1:0") + assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + + +# ============================================================================= +# Normalize Bedrock Profile ID (edge cases) +# ============================================================================= + + +class TestNormalizeBedrockProfileId: + """Test normalization of various Bedrock profile ID formats.""" + + def test_eu_prefixed(self): + assert _normalize_bedrock_profile_id("eu.anthropic.claude-sonnet-4-20250514-v1:0") == ( + "claude-sonnet-4-20250514" + ) + + def test_apac_prefixed(self): + assert _normalize_bedrock_profile_id("apac.anthropic.claude-3-5-sonnet-20241022-v2:0") == ( + "claude-3-5-sonnet-20241022" + ) + + def test_us_prefixed(self): + assert _normalize_bedrock_profile_id("us.anthropic.claude-opus-4-20250514-v1:0") == ( + "claude-opus-4-20250514" + ) + + def test_no_region_prefix(self): + assert _normalize_bedrock_profile_id("anthropic.claude-3-haiku-20240307-v1:0") == ( + "claude-3-haiku-20240307" + ) + + def test_with_bedrock_slash_prefix(self): + assert _normalize_bedrock_profile_id( + "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0" + ) == "claude-sonnet-4-20250514" + + def test_non_claude_returns_none(self): + assert _normalize_bedrock_profile_id("eu.meta.llama-3-70b-v1:0") is None + + def test_already_normalized(self): + assert _normalize_bedrock_profile_id("claude-sonnet-4-20250514") == ( + "claude-sonnet-4-20250514" + )