akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(bedrock): add EU/AP region support with graceful fallback

Browse source 
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.
This commit is contained in:
Yitong Li 2026-03-24 00:13:00 +08:00
parent acddeff384
commit b0ece04b9b
2 changed files with 472 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

148
headroom/backends/litellm.py
View file

@ -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

352
tests/test_bedrock_region.py Normal file
View file

@ -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"
)