fix(bedrock): honor BEDROCK_MANTLE_API_BASE on bedrock/mantle messages and chat URLs

This commit is contained in:
mateo-berri 2026-09-02 11:08:45 -07:00
parent ffc0a8e428
commit dfaf235234
2 changed files with 46 additions and 3 deletions

View file

@ -758,12 +758,13 @@ 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.
"""
override: Final = api_base or aws_bedrock_runtime_endpoint
override: Final = api_base or aws_bedrock_runtime_endpoint or get_secret_str("BEDROCK_MANTLE_API_BASE")
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,42 @@ 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])
def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls):
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT)
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(