diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index cbed2232be5..93306025b02 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, List, Optional from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.bedrock.common_utils import build_mantle_messages_url from litellm.types.llms.openai import AllMessageValues if TYPE_CHECKING: @@ -21,10 +22,6 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -MANTLE_ENDPOINT_TEMPLATE = ( - "https://bedrock-mantle.{region}.api.aws/anthropic/v1/messages" -) - class AmazonMantleConfig(AmazonAnthropicClaudeConfig): """ @@ -46,7 +43,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): stream: Optional[bool] = None, ) -> str: region = self._get_aws_region_name(optional_params=optional_params, model=model) - return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + return build_mantle_messages_url( + api_base=api_base, + aws_bedrock_runtime_endpoint=optional_params.get( + "aws_bedrock_runtime_endpoint" + ), + region=region, + ) def validate_environment( self, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9f58e5c0f1c..5e97394f459 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -622,6 +622,31 @@ def strip_bedrock_throughput_suffix(model: str) -> str: return model +MANTLE_MESSAGES_PATH = "/anthropic/v1/messages" + + +def build_mantle_messages_url( + api_base: Optional[str], + aws_bedrock_runtime_endpoint: Optional[str], + region: str, +) -> str: + """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. + 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 = api_base or aws_bedrock_runtime_endpoint + if override: + base = override.rstrip("/") + if base.endswith(MANTLE_MESSAGES_PATH): + return base + return f"{base}{MANTLE_MESSAGES_PATH}" + return f"https://bedrock-mantle.{region}.api.aws{MANTLE_MESSAGES_PATH}" + + def get_bedrock_base_model(model: str) -> str: """ Get the base model from the given model name. diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index 900d9aa97d8..94e7f90b719 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -8,6 +8,7 @@ stripping that are specific to the bedrock-mantle endpoint. from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from litellm.llms.bedrock.common_utils import build_mantle_messages_url from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) @@ -20,10 +21,6 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -MANTLE_ENDPOINT_TEMPLATE = ( - "https://bedrock-mantle.{region}.api.aws/anthropic/v1/messages" -) - class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): """ @@ -43,7 +40,13 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): stream: Optional[bool] = None, ) -> str: region = self._get_aws_region_name(optional_params=optional_params, model=model) - return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + return build_mantle_messages_url( + api_base=api_base, + aws_bedrock_runtime_endpoint=optional_params.get( + "aws_bedrock_runtime_endpoint" + ), + region=region, + ) def validate_anthropic_messages_environment( self, diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index bbefdd621f0..f7f8f582abc 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -125,6 +125,74 @@ def test_mantle_messages_url_construction(): assert url == "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages" +_VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com" + + +def test_mantle_chat_url_honors_api_base_host(): + config = AmazonMantleConfig() + url = config.get_complete_url( + api_base=_VPC_ENDPOINT, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + +def test_mantle_chat_url_honors_api_base_full_path_without_duplication(): + config = AmazonMantleConfig() + full = f"{_VPC_ENDPOINT}/anthropic/v1/messages" + url = config.get_complete_url( + api_base=full, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == full + + +def test_mantle_messages_url_honors_api_base_host(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=_VPC_ENDPOINT, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + assert "api.aws" not in url + + +def test_mantle_messages_url_honors_api_base_with_trailing_slash(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=f"{_VPC_ENDPOINT}/", + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + +def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={ + "aws_region_name": "us-gov-west-1", + "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT, + }, + 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( @@ -247,3 +315,35 @@ async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body() assert requests[0]["path"] == "/anthropic/v1/messages" assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" assert "aws_bedrock_project_id" not in requests[0]["body"] + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_routes_to_vpc_api_base(): + import litellm + + urls = [] + + async def mock_post(self, url, data=None, headers=None, **kwargs): + urls.append(str(url)) + return _anthropic_response(str(url)) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + api_base=_VPC_ENDPOINT, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-gov-west-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert len(urls) == 1 + assert urls[0] == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + assert "api.aws" not in urls[0]