Merge pull request #39364 from BerriAI/litellm_fix_bedrock_mantle_messages_env_api_base

fix(bedrock): honor BEDROCK_MANTLE_API_BASE on bedrock/mantle messages and chat URLs
This commit is contained in:
Mateo Wang 2026-09-02 16:46:44 -07:00 committed by GitHub
commit 4286be8d5b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 64 additions and 4 deletions

View file

@ -748,6 +748,15 @@ def strip_bedrock_throughput_suffix(model: str) -> str:
MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages"
_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1")
def _mantle_api_base_from_env() -> str | None:
env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE")
if env_base is None:
return None
base: Final = env_base.rstrip("/")
return next((base[: -len(suffix)] for suffix in _MANTLE_OPENAI_BASE_SUFFIXES if base.endswith(suffix)), base)
def build_mantle_messages_url(
@ -758,12 +767,15 @@ def build_mantle_messages_url(
"""Build the bedrock-mantle Anthropic /messages URL.
Honors an explicit endpoint override (``api_base``, then
``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle
endpoints are reachable; otherwise falls back to the public regional host.
``aws_bedrock_runtime_endpoint``, then ``BEDROCK_MANTLE_API_BASE``) so
private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise
falls back to the public regional host.
The mantle messages path is appended unless the override already carries it,
so callers can pass either the host or the full messages URL.
so callers can pass either the host or the full messages URL. The env var is
shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to
carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first.
"""
override: Final = api_base or aws_bedrock_runtime_endpoint
override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env()
if override:
base: Final = override.rstrip("/")
if base.endswith(MANTLE_MESSAGES_PATH):

View file

@ -128,6 +128,12 @@ def test_mantle_messages_url_construction():
_VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com"
@pytest.fixture(autouse=True)
def no_ambient_mantle_api_base(monkeypatch):
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
def test_mantle_chat_url_honors_api_base_host():
config = AmazonMantleConfig()
url = config.get_complete_url(
@ -193,6 +199,48 @@ def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint():
assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages"
_ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com"
@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig])
@pytest.mark.parametrize(
"env_value",
[_ENV_ENDPOINT, f"{_ENV_ENDPOINT}/", f"{_ENV_ENDPOINT}/v1", f"{_ENV_ENDPOINT}/openai/v1"],
)
def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls, env_value):
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", env_value)
url = config_cls().get_complete_url(
api_base=None,
api_key=None,
model="mantle/anthropic.claude-mythos-preview",
optional_params={"aws_region_name": "us-east-1"},
litellm_params={},
)
assert url == f"{_ENV_ENDPOINT}/anthropic/v1/messages"
@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig])
@pytest.mark.parametrize(
("api_base", "optional_params"),
[
(_VPC_ENDPOINT, {"aws_region_name": "us-gov-west-1"}),
(None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}),
],
)
def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env(
monkeypatch, config_cls, api_base, optional_params
):
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT)
url = config_cls().get_complete_url(
api_base=api_base,
api_key=None,
model="mantle/anthropic.claude-mythos-preview",
optional_params=optional_params,
litellm_params={},
)
assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages"
def test_mantle_transform_request_strips_prefix_and_adds_model():
config = AmazonMantleConfig()
request = config.transform_request(