Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions litellm/llms/bedrock/chat/mantle/transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,30 @@ def get_complete_url(
region=region,
)

def sign_request(
self,
headers: dict, # mutable-ok: matches parent class signature
optional_params: dict, # mutable-ok: matches parent class signature
request_data: dict, # mutable-ok: matches parent class signature
api_base: str,
api_key: str | None = None,
model: str | None = None,
stream: bool | None = None,
fake_stream: bool | None = None,
) -> tuple[dict, bytes | None]: # mutable-ok: matches parent class signature
"""Sign request with bedrock-mantle service name instead of bedrock."""
return self._sign_request(
service_name="bedrock-mantle",
headers=headers,
optional_params=optional_params,
request_data=request_data,
api_base=api_base,
api_key=api_key,
model=model,
stream=stream,
fake_stream=fake_stream,
)

def validate_environment(
self,
headers: dict,
Expand Down
4 changes: 2 additions & 2 deletions litellm/llms/bedrock_mantle/common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
Mantle authenticates with a Bearer token when one is available
(litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the standard
AWS_BEARER_TOKEN_BEDROCK); otherwise it falls back to AWS SigV4 (service
"bedrock") over the standard credential chain (IAM role / access key / profile /
"bedrock-mantle") over the standard credential chain (IAM role / access key / profile /
web identity). The Chat Completions and Responses backends share this behaviour
through BedrockMantleAuthMixin so the two paths can never drift apart.

Expand Down Expand Up @@ -83,7 +83,7 @@ def sign_request(
headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
try:
return self._aws_signer._sign_request(
service_name="bedrock",
service_name="bedrock-mantle",
headers=headers,
optional_params=optional_params,
request_data=request_data,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1211,7 +1211,7 @@ def test_access_key_produces_sigv4_headers(self, monkeypatch):
)
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "Credential=AKIAEXAMPLE/" in headers["Authorization"]
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "/us-east-2/bedrock-mantle/aws4_request" in headers["Authorization"]
assert "X-Amz-Date" in headers
assert headers["X-Amz-Security-Token"] == "session-token-test"
assert signed_body == b'{"input": "hi"}'
Expand Down Expand Up @@ -1250,7 +1250,7 @@ def test_assume_role_path_produces_sigv4_headers(self, monkeypatch):
assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role"
assert call["aws_session_name"] == "litellm-test"
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "/us-east-2/bedrock-mantle/aws4_request" in headers["Authorization"]

def test_signed_body_matches_final_data_after_normalize(self, monkeypatch):
"""Core regression: the signed bytes must equal the bytes actually sent.
Expand Down Expand Up @@ -1300,7 +1300,7 @@ def test_region_comes_from_optional_params(self, monkeypatch):
api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses",
api_key=None,
)
assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
assert "/eu-west-1/bedrock-mantle/aws4_request" in headers["Authorization"]

def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch):
"""Adversarial-review regression: a caller-supplied aws_region_name (no region
Expand Down Expand Up @@ -1334,7 +1334,7 @@ def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch
api_base=url,
api_key=None,
)
assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"]
assert "/ap-southeast-2/bedrock-mantle/aws4_request" in headers["Authorization"]

def test_injected_default_region_base_does_not_override_aws_region_name(
self, monkeypatch
Expand Down Expand Up @@ -1372,7 +1372,7 @@ def test_injected_default_region_base_does_not_override_aws_region_name(
api_base=url,
api_key=None,
)
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "/us-east-2/bedrock-mantle/aws4_request" in headers["Authorization"]
assert "us-east-1" not in headers["Authorization"]

def test_custom_proxy_host_is_preserved(self, monkeypatch):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -329,7 +329,7 @@ def test_no_bearer_signs_with_sigv4(self, monkeypatch):

assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "Credential=AKIAEXAMPLE/" in headers["Authorization"]
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert "/us-east-2/bedrock-mantle/aws4_request" in headers["Authorization"]
assert headers["X-Amz-Security-Token"] == "session-token-test"
assert json.loads(signed_body) == {
"model": "openai.gpt-oss-120b",
Expand Down Expand Up @@ -364,7 +364,7 @@ def test_sigv4_region_resolved_from_api_base_host(self, monkeypatch):
api_key=None,
)

assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
assert "/eu-west-1/bedrock-mantle/aws4_request" in headers["Authorization"]

def test_sigv4_scope_matches_api_base_when_aws_region_name_disagrees(
self, monkeypatch
Expand Down Expand Up @@ -399,8 +399,8 @@ def test_sigv4_scope_matches_api_base_when_aws_region_name_disagrees(
api_key=None,
)

assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
assert "/us-west-2/bedrock/aws4_request" not in headers["Authorization"]
assert "/eu-west-1/bedrock-mantle/aws4_request" in headers["Authorization"]
assert "/us-west-2/bedrock-mantle/aws4_request" not in headers["Authorization"]

def test_no_bearer_and_no_credentials_raises_value_error(self, monkeypatch):
from unittest.mock import MagicMock
Expand Down Expand Up @@ -486,7 +486,7 @@ def mock_post(self, url, data=None, headers=None, **kwargs):
assert len(requests) == 1
authorization = requests[0]["headers"]["Authorization"]
assert authorization.startswith("AWS4-HMAC-SHA256")
assert "/us-east-2/bedrock/aws4_request" in authorization
assert "/us-east-2/bedrock-mantle/aws4_request" in authorization
assert requests[0]["url"].startswith("https://bedrock-mantle.us-east-2.api.aws")


Expand Down
Loading