diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1e634ced29b..1f00bf7792e 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -49,11 +49,16 @@ SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz class Boto3CredentialsInfo(BaseModel): - credentials: Credentials + credentials: Credentials | None aws_region_name: str aws_bedrock_runtime_endpoint: str | 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 @@ -1388,7 +1393,7 @@ class BaseAWSLLM: return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" def _get_boto_credentials_from_optional_params( - self, optional_params: dict, model: str | None = None + self, optional_params: dict, model: str | None = None, bearer_token: str | None = None ) -> Boto3CredentialsInfo: """ Get boto3 credentials from optional params @@ -1420,17 +1425,21 @@ class BaseAWSLLM: ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) - 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 Boto3CredentialsInfo( @@ -1451,14 +1460,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 +1559,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}" diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7d5f99ca893..a75124325ae 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -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 ### diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index c34ca7750e2..a7b74f3752a 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -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 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 @@ -45,11 +45,8 @@ class BedrockEmbedding(BaseAWSLLM): 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 +75,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 +234,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 +302,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 +384,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 @@ -595,8 +598,11 @@ class BedrockEmbedding(BaseAWSLLM): try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest + from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + if credentials is None: + raise NoCredentialsError() # Create AWSRequest with GET method and encoded URL request: Final = AWSRequest( diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 9d8631c7c26..5c517f2049c 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -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) diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index 6fac14a0dc3..c78e3c147cb 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 30526d30dc5..7f2616c2fb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 2c7f1f8389d..acdb83094e6 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -535,6 +535,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest + from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") optional_params = optional_params or {} @@ -582,10 +583,14 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): "X-Amz-Target": f"secretsmanager.{action}", } + credentials: Final = boto3_credentials_info.credentials + if credentials is None: + raise NoCredentialsError() + # Sign request request: Final = AWSRequest(method="POST", url=endpoint_url, data=body, headers=headers) SigV4Auth( - boto3_credentials_info.credentials, + credentials, "secretsmanager", boto3_credentials_info.aws_region_name, ).add_auth(request) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 21e3239f623..c4d6896b17b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -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" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 08d01127eba..50f8bbcf584 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -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" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py index 7c36b2aa75f..0b11a66c100 100644 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -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" diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py index 020b8df1276..58411a9ae18 100644 --- a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py +++ b/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py @@ -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" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 953e3de1519..14e124981c9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5792,3 +5792,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" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index d842a1ee5f9..f4af77d2e40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -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"