mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(bedrock-mantle): honor api_base for VPC endpoint routing on bedrock/mantle/... (#31141)
The bedrock/mantle chat and messages paths hardcoded the public Mantle host and ignored api_base, so private VPC/VPCE/GovCloud endpoints could not be used. Route URL construction through a shared helper that prefers api_base and aws_bedrock_runtime_endpoint before falling back to the regional public host. Co-authored-by: Shivam Rawat <shivamrawat@Shivams-MacBook-Pro.local> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
aab732d874
commit
8dfe702e5b
4 changed files with 141 additions and 10 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue