diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1e634ced29b..c3da992a904 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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}" 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/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 91c3a363c31..04c6ec86a13 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -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, ) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index c34ca7750e2..5fb86d476f4 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, 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 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/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py new file mode 100644 index 00000000000..531f334e460 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py @@ -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 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 479d1f2d4b2..1e44d29b610 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 @@ -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" 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" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index b99dcb062b4..18aa43f7d1c 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -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