scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge pull request #39411 from BerriAI/litellm_bedrock_bearer_skip_sigv4_chain

Browse source 
fix(bedrock): skip the SigV4 credential chain when a bearer token is configured
This commit is contained in:
Mateo Wang 2026-09-03 14:36:27 -07:00 • committed by GitHub
commit f7691a3d85
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 346 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

63
litellm/llms/bedrock/base_aws_llm.py
View file

@ -7,7 +7,7 @@ import urllib.parse
from collections.abc import Callable
from datetime import datetime
from threading import Lock
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload
import httpx
from pydantic import BaseModel, ValidationError
@ -48,12 +48,24 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile(
SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"})
class Boto3CredentialsInfo(BaseModel):
credentials: Credentials
class BedrockRequestTarget(BaseModel):
aws_region_name: str
aws_bedrock_runtime_endpoint: str | None
class Boto3CredentialsInfo(BedrockRequestTarget):
credentials: Credentials
class BearerRequestTarget(BedrockRequestTarget):
credentials: None = None
def bedrock_bearer_token(api_key: str | None) -> str | None:
token: Final = api_key if api_key is not None else get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
return token or None
class _WebIdentityTokenClaims(BaseModel):
aud: str | list[str] | None = None
iss: str | None = None
@ -1387,9 +1399,26 @@ class BaseAWSLLM:
else:
return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}"
@overload
def _get_boto_credentials_from_optional_params(
self, optional_params: dict, model: str | None = None
) -> Boto3CredentialsInfo:
self,
optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place
model: str | None = None,
bearer_token: None = None,
) -> Boto3CredentialsInfo: ...
@overload
def _get_boto_credentials_from_optional_params(
self,
optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place
model: str | None = None,
*,
bearer_token: str,
) -> BearerRequestTarget: ...
def _get_boto_credentials_from_optional_params(
self, optional_params: dict, model: str | None = None, bearer_token: str | None = None
) -> Boto3CredentialsInfo | BearerRequestTarget:
"""
Get boto3 credentials from optional params
@ -1420,6 +1449,12 @@ class BaseAWSLLM:
) # https://bedrock-runtime.{region_name}.amazonaws.com
aws_external_id: Final = optional_params.pop("aws_external_id", None)
if bearer_token is not None:
return BearerRequestTarget(
aws_region_name=aws_region_name,
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
)
credentials: Final[Credentials] = self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
@ -1432,7 +1467,6 @@ class BaseAWSLLM:
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
return Boto3CredentialsInfo(
credentials=credentials,
aws_region_name=aws_region_name,
@ -1451,14 +1485,9 @@ class BaseAWSLLM:
api_key: str | None = None,
supports_bearer_token: bool = True,
) -> AWSPreparedRequest:
if not supports_bearer_token:
aws_bearer_token: str | None = None
elif api_key is not None:
aws_bearer_token = api_key
else:
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
aws_bearer_token: Final = bedrock_bearer_token(api_key) if supports_bearer_token else None
if aws_bearer_token:
if aws_bearer_token is not None:
try:
from botocore.awsrequest import AWSRequest
except ImportError:
@ -1555,13 +1584,9 @@ class BaseAWSLLM:
Returns:
Tuple[dict, Optional[str]]: A tuple containing the headers and the json str body of the request
"""
if api_key is not None:
aws_bearer_token: str | None = api_key
else:
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
aws_bearer_token: Final = bedrock_bearer_token(api_key)
# If aws bearer token is set, use it directly in the header
if aws_bearer_token:
if aws_bearer_token is not None:
headers = headers or {}
headers["Content-Type"] = "application/json"
headers["Authorization"] = f"Bearer {aws_bearer_token}"

28
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
from ..base_aws_llm import BaseAWSLLM, Credentials
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
from ..common_utils import BedrockError, _get_all_bedrock_regions
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
@ -349,17 +349,21 @@ class BedrockConverseLLM(BaseAWSLLM):
litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls
credentials: Final[Credentials | None] = self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
credentials: Final[Credentials | None] = (
None
if bedrock_bearer_token(api_key) is not None
else self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
)
### SET RUNTIME ENDPOINT ###

8
litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
View file

@ -149,19 +149,15 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
- Temperature and parameter validation
"""
# Filter out AWS credentials using the existing method from BaseAWSLLM
self._get_boto_credentials_from_optional_params(optional_params, model)
inference_params: Final = {k: v for k, v in optional_params.items() if k not in self.aws_authentication_params}
# Strip routing prefixes to get the actual model ID
clean_model_id: Final = self._get_model_id(model)
# Use Moonshot's transform_request which handles message transformation
# and tool_choice="required" workaround
return MoonshotChatConfig.transform_request(
self,
model=clean_model_id,
messages=messages,
optional_params=optional_params,
optional_params=inference_params,
litellm_params=litellm_params,
headers=headers,
)

59
litellm/llms/bedrock/embed/embedding.py
View file

@ -6,7 +6,7 @@ import copy
import json
import urllib.parse
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Final, get_args
from typing import TYPE_CHECKING, Final, get_args, overload
import httpx
@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import (
)
from litellm.types.utils import EmbeddingResponse, LlmProviders
from ..base_aws_llm import BaseAWSLLM
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
from ..common_utils import BedrockError
from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
from .amazon_titan_g1_transformation import AmazonTitanG1Config
@ -42,14 +42,25 @@ if TYPE_CHECKING:
class BedrockEmbedding(BaseAWSLLM):
@overload
def _load_credentials(
self,
optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place
bearer_token: None = None,
) -> tuple[Credentials, str]: ...
@overload
def _load_credentials(
self,
optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place
bearer_token: str,
) -> tuple[None, str]: ...
def _load_credentials(
self,
optional_params: dict,
) -> tuple[Any, str]:
try:
from botocore.credentials import Credentials
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
bearer_token: str | None = None,
) -> tuple[Credentials | None, str]:
## CREDENTIALS ##
# pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them
aws_secret_access_key: Final = optional_params.pop("aws_secret_access_key", None)
@ -78,17 +89,21 @@ class BedrockEmbedding(BaseAWSLLM):
if aws_region_name is None:
aws_region_name = "us-west-2"
credentials: Final[Credentials] = self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
credentials: Final[Credentials | None] = (
None
if bearer_token is not None
else self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
)
return credentials, aws_region_name
@ -233,7 +248,7 @@ class BedrockEmbedding(BaseAWSLLM):
client: HTTPHandler | None,
timeout: float | httpx.Timeout | None,
batch_data: list[dict],
credentials: Any,
credentials: Credentials | None,
extra_headers: dict | None,
endpoint_url: str,
aws_region_name: str,
@ -301,7 +316,7 @@ class BedrockEmbedding(BaseAWSLLM):
client: AsyncHTTPHandler | None,
timeout: float | httpx.Timeout | None,
batch_data: list[dict],
credentials: Any,
credentials: Credentials | None,
extra_headers: dict | None,
endpoint_url: str,
aws_region_name: str,
@ -383,7 +398,9 @@ class BedrockEmbedding(BaseAWSLLM):
litellm_params: dict,
api_key: str | None = None,
) -> EmbeddingResponse:
credentials, aws_region_name = self._load_credentials(optional_params)
credentials, aws_region_name = self._load_credentials(
optional_params, bearer_token=bedrock_bearer_token(api_key)
)
### TRANSFORMATION ###
unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed

6
litellm/llms/bedrock/image_edit/handler.py
View file

@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.types.utils import ImageResponse
from ..base_aws_llm import BaseAWSLLM
from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token
from ..common_utils import BedrockError
if TYPE_CHECKING:
@ -198,7 +198,9 @@ class BedrockImageEdit(BaseAWSLLM):
Returns:
BedrockImageEditPreparedRequest: The prepared request object
"""
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(
optional_params, model, bearer_token=bedrock_bearer_token(api_key)
)
# Use the existing ARN-aware provider detection method
bedrock_provider: Final = self.get_bedrock_invoke_provider(model)

6
litellm/llms/bedrock/image_generation/image_handler.py
View file

@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.types.utils import ImageResponse
from ..base_aws_llm import BaseAWSLLM
from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token
from ..common_utils import BedrockError
if TYPE_CHECKING:
@ -220,7 +220,9 @@ class BedrockImageGeneration(BaseAWSLLM):
prepped (httpx.Request): The prepared request object
body (bytes): The request body
"""
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(
optional_params, model, bearer_token=bedrock_bearer_token(api_key)
)
# Use the existing ARN-aware provider detection method
bedrock_provider: Final = self.get_bedrock_invoke_provider(model)

53
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -41,7 +41,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -56,7 +56,6 @@ from litellm.proxy.guardrails.anthropic_sse import (
is_raw_sse_stream,
model_response_text,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import (
BedrockChecksConfigModel,
BedrockGuardrailStreamingParams,
@ -713,9 +712,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# logic becomes shared across providers.
#### CALL HOOKS - proxy only ####
def _load_credentials(
self,
):
def _load_credentials(self, bearer_token: str | None = None):
try:
from botocore.credentials import Credentials
except ImportError:
@ -737,17 +734,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
aws_region_name=aws_region_name,
)
credentials: Final[Credentials] = self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
credentials: Final[Credentials | None] = (
None
if bearer_token is not None
else self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
)
return credentials, aws_region_name
@ -779,13 +780,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
proxy_endpoint_url = f"{proxy_endpoint_url}{request_path}"
encoded_data: Final = json.dumps(data).encode("utf-8")
# first check api-key, if none, fall back to sigV4
if api_key is not None:
aws_bearer_token: str | None = api_key
else:
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
aws_bearer_token: Final = bedrock_bearer_token(api_key)
if aws_bearer_token:
if aws_bearer_token is not None:
try:
from botocore.awsrequest import AWSRequest
except ImportError:
@ -916,7 +913,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
source,
)
return BedrockGuardrailResponse()
credentials, aws_region_name = self._load_credentials()
credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
allow_chunking: Final = not self._content_uses_contextual_grounding(content)
completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
@ -958,7 +955,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self,
content: Sequence[BedrockContentItem],
base_request_data: Mapping[str, object],
credentials: "Credentials",
credentials: "Credentials | None",
aws_region_name: str,
api_key: str | None,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
@ -1096,7 +1093,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self,
content: Sequence[BedrockContentItem],
base_request_data: Mapping[str, object],
credentials: "Credentials",
credentials: "Credentials | None",
aws_region_name: str,
api_key: str | None,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
@ -1146,7 +1143,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self,
content: Sequence[BedrockContentItem],
base_request_data: Mapping[str, object],
credentials: "Credentials",
credentials: "Credentials | None",
aws_region_name: str,
api_key: str | None,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
@ -1873,9 +1870,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does.
return BedrockGuardrailResponse()
credentials, aws_region_name = self._load_credentials()
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
prepared_request: Final = self._prepare_request(
credentials=credentials,

66
tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py Normal file
View file

@ -0,0 +1,66 @@
import pytest
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
AWS_AUTH_PARAMS = {
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "secret",
"aws_session_token": "token",
"aws_region_name": "us-west-2",
"aws_session_name": "session",
"aws_role_name": "arn:aws:iam::000000000000:role/example",
"aws_web_identity_token": "web-identity",
"aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com",
"aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com",
"aws_external_id": "external",
}
def test_transform_request_never_resolves_aws_credentials():
"""A broken credential chain must not stop the request body from being built."""
config = AmazonMoonshotConfig()
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[{"role": "user", "content": "Hello"}],
optional_params={"aws_profile_name": "litellm-profile-that-does-not-exist", "max_tokens": 16},
litellm_params={},
headers={},
)
assert transformed["model"] == "moonshot.kimi-k2-thinking"
assert transformed["max_tokens"] == 16
assert "aws_profile_name" not in transformed
@pytest.mark.parametrize("aws_param", sorted(AWS_AUTH_PARAMS))
def test_transform_request_keeps_aws_params_out_of_the_body(aws_param: str):
config = AmazonMoonshotConfig()
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[{"role": "user", "content": "Hello"}],
optional_params={aws_param: AWS_AUTH_PARAMS[aws_param]},
litellm_params={},
headers={},
)
assert aws_param not in transformed
def test_transform_request_leaves_the_caller_aws_params_in_place_for_signing():
"""sign_request reads the aws_* keys off optional_params after transform_request runs."""
config = AmazonMoonshotConfig()
optional_params = dict(AWS_AUTH_PARAMS)
config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[{"role": "user", "content": "Hello"}],
optional_params=optional_params,
litellm_params={},
headers={},
)
assert optional_params == AWS_AUTH_PARAMS

24
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -513,3 +513,27 @@ def test_the_rust_opt_in_needs_no_sigv4_principal():
assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys()
assert params["aws_region_name"] == "us-east-1"
assert seen["call"][0]["api_key"] == "bedrock-bearer-token"
@pytest.mark.parametrize("configured_through", ["env_var", "api_key"])
def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, configured_through):
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
raises; a bearer-token deployment must still serve the request, since the
bearer token alone signs it."""
if configured_through == "env_var":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
else:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
client = _sync_client_returning_converse_response()
response = BedrockConverseLLM().completion(
**_completion_kwargs(
optional_params={"maxTokens": 16, "aws_profile_name": "litellm-no-such-aws-profile"},
litellm_params={},
client=client,
api_key="bedrock-bearer-token" if configured_through == "api_key" else None,
)
)
assert response.choices[0].message.content == "hi"
assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token"

26
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -1033,3 +1033,29 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch):
assert credentials.token == "assumed-session-token"
assert aws_region_name == "us-east-1"
assert "aws_external_id" not in optional_params
def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
raises; a bearer-token deployment must still serve the request, since the
bearer token alone signs it."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
client = HTTPHandler()
with patch.object(client, "post") as mock_post:
mock_response = Mock()
mock_response.status_code = 200
mock_response.text = json.dumps(titan_embedding_response)
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
response = litellm.embedding(
model="bedrock/amazon.titan-embed-text-v1",
input=test_input,
client=client,
aws_region_name="us-west-2",
aws_profile_name="litellm-no-such-aws-profile",
)
assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"

21
tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py
View file

@ -135,3 +135,24 @@ class TestBedrockImageGeneration:
assert response is not None
assert len(response.data) > 0
mock_bedrock_image_gen.assert_called_once()
def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
raises; a bearer-token deployment must still sign the request with the
bearer token alone."""
from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
request = BedrockImageGeneration()._prepare_request(
model="amazon.nova-canvas-v1:0",
prompt="A cute baby sea otter",
optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"},
api_base=None,
extra_headers=None,
api_key=None,
logging_obj=Mock(),
)
assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345"

21
tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py
View file

@ -3,6 +3,7 @@
import base64
import io
from typing import cast
from unittest.mock import Mock, patch
import httpx
import pytest
@ -655,3 +656,23 @@ def test_transform_response_empty_images_without_error_raises():
raw_response=resp,
logging_obj=None, # type: ignore[arg-type]
)
def test_prepare_request_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
raises; a bearer-token deployment must still sign the request with the
bearer token alone."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
request = BedrockImageEdit()._prepare_request(
model="amazon.nova-canvas-v1:0",
image=[io.BytesIO(b"fake-png")],
prompt="make it warmer",
optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"},
api_base=None,
extra_headers=None,
logging_obj=Mock(),
api_key=None,
)
assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345"

22
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -5784,3 +5784,25 @@ async def test_apply_guardrail_debug_log_masks_signed_request_headers():
assert header_lines, "expected the signed-request debug line to be logged"
assert any("X-Amz-Security-Token" in message for message in header_lines)
assert all(session_token not in message for message in rendered_messages)
@pytest.mark.asyncio
async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
"""The guardrail's AWS profile does not exist, so resolving SigV4 credentials
raises; with a bearer token configured the guardrail must still run, since
the bearer token alone signs the request."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
aws_profile_name="litellm-no-such-aws-profile",
)
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"action": "NONE", "assessments": []}
with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response) as mock_post:
response = await guardrail.make_bedrock_api_request(source="INPUT", messages=[{"role": "user", "content": "hello"}])
assert response["action"] == "NONE"
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"

28
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py
View file

@ -833,3 +833,31 @@ async def test_many_blocks_scanned_at_request_level_and_can_block():
sent_texts = [c["text"] for m in body_messages for c in m["content"]]
assert sent_texts == [f"b{i}" for i in range(25)]
assert all(len(m["content"]) <= 10 for m in body_messages)
@pytest.mark.asyncio
async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
"""Same bearer-token rule as ApplyGuardrail: the guardrail's AWS profile does
not exist, yet the InvokeGuardrailChecks call still goes out on the bearer
token and its verdict is enforced."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
g = BedrockGuardrail(
checks=CONTENT_FILTER_CHECKS,
content_filter_threshold=0.5,
aws_profile_name="litellm-no-such-aws-profile",
)
payload = {"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.8}]}}}
post = AsyncMock(return_value=_mock_http_response(200, payload))
with patch.object(g.async_handler, "post", new=post):
with pytest.raises(HTTPException) as exc:
await g.make_bedrock_api_request(
source="INPUT",
messages=[{"role": "user", "content": "hi"}],
request_data={"messages": []},
)
assert exc.value.detail["bedrock_guardrail_checks"] == [
{"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8}
]
assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"

21
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -802,7 +802,7 @@ async def test_bedrock_guardrail_prepare_request_with_api_key():
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_without_api_key():
async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch):
"""Test _prepare_request method falls back to SigV4 when no api_key is provided"""
from unittest.mock import Mock, patch
@ -820,18 +820,13 @@ async def test_bedrock_guardrail_prepare_request_without_api_key():
# Test data without api_key
test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
with (
patch(
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str"
) as mock_get_secret,
patch("botocore.auth.SigV4Auth") as mock_sigv4_auth,
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
):
# Mock no AWS_BEARER_TOKEN_BEDROCK
mock_get_secret.return_value = None
# Mock SigV4Auth
mock_sigv4_instance = Mock()
mock_sigv4_auth.return_value = mock_sigv4_instance
@ -857,7 +852,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key():
@pytest.mark.asyncio
async def test_bedrock_guardrail_prepare_request_with_bearer_token_env():
async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(monkeypatch):
"""Test _prepare_request method uses Bearer token from environment when available"""
from unittest.mock import Mock, patch
@ -875,15 +870,9 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env():
# Test data without api_key
test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-456")
with (
patch(
"litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str"
) as mock_get_secret,
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
):
mock_get_secret.return_value = "env-bearer-token-456"
with patch("botocore.awsrequest.AWSRequest") as mock_aws_request:
mock_request_instance = Mock()
mock_request_instance.prepare.return_value = Mock()
mock_aws_request.return_value = mock_request_instance