From d5060b361192563b1d56d78576dbb5eef0093381 Mon Sep 17 00:00:00 2001 From: jjj-n <1357598741@qq.com> Date: Fri, 28 Aug 2026 11:16:12 +0800 Subject: [PATCH 1/2] fix(bedrock): skip AWS credentials for bearer auth --- litellm/llms/bedrock/base_aws_llm.py | 71 ++++++++++++------- litellm/llms/bedrock/chat/converse_handler.py | 54 ++++++++------ .../amazon_moonshot_transformation.py | 7 +- litellm/llms/bedrock/embed/embedding.py | 40 +++++++---- litellm/llms/bedrock/image_edit/handler.py | 7 +- .../bedrock/image_generation/image_handler.py | 7 +- .../llm_translation/test_bedrock_moonshot.py | 22 ++++++ .../chat/test_bedrock_converse_handler.py | 19 +++++ .../llms/bedrock/test_base_aws_llm.py | 59 ++++++++++++++- 9 files changed, 221 insertions(+), 65 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 4332848e545..b77c6e8e739 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -48,7 +48,7 @@ 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 @@ -1372,8 +1372,24 @@ class BaseAWSLLM: else: return f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" + def _get_aws_bearer_token( + self, + api_key: str | None = None, + supports_bearer_token: bool = True, + ) -> str | None: + if not supports_bearer_token: + return None + if api_key is not None: + return api_key + return get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + def _get_boto_credentials_from_optional_params( - self, optional_params: dict, model: str | None = None + self, + optional_params: dict, + model: str | None = None, + *, + api_key: str | None = None, + supports_bearer_token: bool = False, ) -> Boto3CredentialsInfo: """ Get boto3 credentials from optional params @@ -1382,7 +1398,8 @@ class BaseAWSLLM: optional_params (dict): Optional parameters for the model call Returns: - Credentials: Boto3 credentials object + Boto3CredentialsInfo: Resolved request metadata. Credentials are None + when bearer-token authentication is active. """ try: from botocore.credentials import Credentials @@ -1405,17 +1422,24 @@ 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 self._get_aws_bearer_token( + api_key=api_key, + supports_bearer_token=supports_bearer_token, + ) + 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( @@ -1427,7 +1451,7 @@ class BaseAWSLLM: @tracer.wrap() def get_request_headers( self, - credentials: Credentials, + credentials: Credentials | None, aws_region_name: str, extra_headers: dict | None, endpoint_url: str, @@ -1436,12 +1460,10 @@ 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 = self._get_aws_bearer_token( + api_key=api_key, + supports_bearer_token=supports_bearer_token, + ) if aws_bearer_token: try: @@ -1451,6 +1473,8 @@ class BaseAWSLLM: headers["Authorization"] = f"Bearer {aws_bearer_token}" request = AWSRequest(method="POST", url=endpoint_url, data=data, headers=headers) else: + if credentials is None: + raise ValueError("AWS credentials are required for SigV4 authentication") try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest @@ -1536,10 +1560,7 @@ 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 = self._get_aws_bearer_token(api_key=api_key) # If aws bearer token is set, use it directly in the header if aws_bearer_token: diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index ca5f1298360..462a1e3ab24 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -95,7 +95,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream, optional_params: dict, litellm_params: dict, - credentials: Credentials, + credentials: Credentials | None, logger_fn=None, headers={}, client: AsyncHTTPHandler | None = None, @@ -167,7 +167,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream, optional_params: dict, litellm_params: dict, - credentials: Credentials, + credentials: Credentials | None, logger_fn=None, headers: dict = {}, client: AsyncHTTPHandler | None = None, @@ -331,17 +331,21 @@ class BedrockConverseLLM(BaseAWSLLM): litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls - 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 self._get_aws_bearer_token(api_key=api_key) + 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 ### @@ -371,16 +375,20 @@ class BedrockConverseLLM(BaseAWSLLM): # resolved so both paths sign as the same principal. rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy **optional_params, - **{ # mutable-ok: merged into its mutable parent above - key: value - for key, value in ( - ("aws_access_key_id", credentials.access_key), - ("aws_secret_access_key", credentials.secret_key), - ("aws_session_token", credentials.token), - ("aws_region_name", aws_region_name), - ) - if value is not None - }, + "aws_region_name": aws_region_name, + **( + { # mutable-ok: merged into its mutable parent above + key: value + for key, value in ( + ("aws_access_key_id", credentials.access_key), + ("aws_secret_access_key", credentials.secret_key), + ("aws_session_token", credentials.token), + ) + if value is not None + } + if credentials is not None + else {} + ), } serves_via_rust: Final = rust_chat_completions_accepts( model=model, 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 a8275f1d35f..871cde9c1bf 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -148,7 +148,12 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): """ # Filter out AWS credentials using the existing method from BaseAWSLLM - self._get_boto_credentials_from_optional_params(optional_params, model) + self._get_boto_credentials_from_optional_params( + optional_params, + model, + api_key=litellm_params.get("api_key"), + supports_bearer_token=True, + ) # Strip routing prefixes to get the actual model ID clean_model_id: Final = self._get_model_id(model) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 082bf7ee2d9..60db0f70eff 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -42,6 +42,9 @@ class BedrockEmbedding(BaseAWSLLM): def _load_credentials( self, optional_params: dict, + *, + api_key: str | None = None, + supports_bearer_token: bool = True, ) -> tuple[Any, str]: try: from botocore.credentials import Credentials @@ -74,16 +77,23 @@ 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, + credentials: Final[Credentials | None] = ( + None + if self._get_aws_bearer_token( + api_key=api_key, + supports_bearer_token=supports_bearer_token, + ) + 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, + ) ) return credentials, aws_region_name @@ -378,7 +388,10 @@ 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, + api_key=api_key, + ) ### TRANSFORMATION ### unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed @@ -568,7 +581,10 @@ class BedrockEmbedding(BaseAWSLLM): """ # Get AWS credentials using the same method as other Bedrock methods - credentials, _ = self._load_credentials(kwargs) + credentials, _ = self._load_credentials( + kwargs, + supports_bearer_token=False, + ) # Get the runtime endpoint endpoint_url, _ = self.get_runtime_endpoint( diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 9d8631c7c26..8f4689835d9 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -198,7 +198,12 @@ 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, + api_key=api_key, + supports_bearer_token=True, + ) # 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..d7abe15b2f0 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -220,7 +220,12 @@ 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, + api_key=api_key, + supports_bearer_token=True, + ) # Use the existing ARN-aware provider detection method bedrock_provider: Final = self.get_bedrock_invoke_provider(model) diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index 3bf047c51a5..038ccf9b82f 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -329,6 +329,28 @@ class TestBedrockMoonshotBasic: # The model ID in the request body should be stripped assert transformed["model"] == "moonshot.kimi-k2-thinking" + def test_transform_request_bearer_token_skips_aws_credentials(self): + from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, + ) + + config = AmazonMoonshotConfig() + with patch.object( + config, + "get_credentials", + side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"), + ) as get_credentials: + transformed = config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"api_key": "bedrock-bearer-token"}, + headers={}, + ) + + get_credentials.assert_not_called() + assert transformed["model"] == "moonshot.kimi-k2-thinking" + class TestBedrockMoonshotReasoningContent: """Tests for reasoning content extraction.""" 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 8e67a7e3438..3a7ed9d806f 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 @@ -135,6 +135,25 @@ def test_the_core_receives_the_credentials_this_handler_already_resolved(): assert params["aws_region_name"] == "us-east-1" +def test_bearer_token_skips_aws_credential_resolution(): + seen = _inject() + + with patch.object( + BedrockConverseLLM, + "get_credentials", + side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"), + ) as get_credentials: + response = BedrockConverseLLM().completion(**_completion_kwargs(api_key="bedrock-bearer-token")) + + get_credentials.assert_not_called() + assert response.choices[0].message.content == "hello from rust" + params = seen["call"][0]["optional_params"] + assert params["aws_region_name"] == "us-east-1" + assert "aws_access_key_id" not in params + assert "aws_secret_access_key" not in params + assert "aws_session_token" not in params + + def test_the_core_receives_the_converse_url_this_handler_already_built(): seen = _inject() _run() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 50e2b53c2b3..60b23defd1e 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -731,10 +731,65 @@ def test_sign_request_with_api_key_bearer_token(): assert result_body == json.dumps(request_data).encode() +@pytest.mark.parametrize( + ("api_key", "env_bearer_token"), + [ + ("deployment-bearer-token", None), + (None, "environment-bearer-token"), + ], +) +def test_bearer_token_skips_boto_credential_resolution( + api_key: str | None, + env_bearer_token: str | None, +): + llm = BaseAWSLLM() + optional_params = {"aws_region_name": "us-east-1"} + + with ( + patch( + "litellm.llms.bedrock.base_aws_llm.get_secret_str", + return_value=env_bearer_token, + ), + patch.object( + llm, + "get_credentials", + side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"), + ) as get_credentials, + ): + credential_info = llm._get_boto_credentials_from_optional_params( + optional_params, + api_key=api_key, + supports_bearer_token=True, + ) + + get_credentials.assert_not_called() + assert credential_info.credentials is None + assert credential_info.aws_region_name == "us-east-1" + + +def test_shared_boto_helper_requires_credentials_by_default(): + llm = BaseAWSLLM() + credentials = Credentials("test_key", "test_secret", "test_token") + + with ( + patch( + "litellm.llms.bedrock.base_aws_llm.get_secret_str", + return_value="environment-bearer-token", + ), + patch.object(llm, "get_credentials", return_value=credentials) as get_credentials, + ): + credential_info = llm._get_boto_credentials_from_optional_params( + {"aws_region_name": "us-east-1"}, + ) + + get_credentials.assert_called_once() + assert credential_info.credentials is credentials + + def test_get_request_headers_with_env_var_bearer_token(): # Setup llm = BaseAWSLLM() - credentials = Credentials("test_key", "test_secret", "test_token") + credentials = None headers = {"Content-Type": "application/json"} headers_dict = headers.copy() @@ -831,7 +886,7 @@ def test_get_request_headers_with_api_key_bearer_token(): """ # Setup llm = BaseAWSLLM() - credentials = Credentials("test_key", "test_secret", "test_token") + credentials = None headers = {"Content-Type": "application/json"} headers_dict = headers.copy() api_key = "test_api_key" From 4de35621dda66fb8157fc07f0835edbe3197acb4 Mon Sep 17 00:00:00 2001 From: jjj-n <1357598741@qq.com> Date: Fri, 28 Aug 2026 13:08:18 +0800 Subject: [PATCH 2/2] fix(bedrock): resolve bearer auth CI failures --- litellm/llms/bedrock/chat/converse_handler.py | 31 +++++++------- .../llm_translation/test_bedrock_moonshot.py | 23 ---------- .../test_amazon_moonshot_transformation.py | 20 +++++++++ .../chat/test_bedrock_converse_handler.py | 16 +++---- .../test_bedrock_async_invoke_embedding.py | 30 ++++++++++++- .../bedrock/embed/test_bedrock_embedding.py | 1 + .../test_amazon_nova_canvas_image_edit.py | 19 +++++++++ .../llms/bedrock/test_base_aws_llm.py | 42 +++++++++---------- 8 files changed, 115 insertions(+), 67 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 462a1e3ab24..90378d33d3a 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -373,22 +373,25 @@ class BedrockConverseLLM(BaseAWSLLM): # before transforming so whichever path runs emits pre_call once, and # hand down the credentials, region and endpoint this handler already # resolved so both paths sign as the same principal. + credential_params: Final = ( + () + if credentials is None + else ( + ("aws_access_key_id", credentials.access_key), + ("aws_secret_access_key", credentials.secret_key), + ("aws_session_token", credentials.token), + ) + ) rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy **optional_params, - "aws_region_name": aws_region_name, - **( - { # mutable-ok: merged into its mutable parent above - key: value - for key, value in ( - ("aws_access_key_id", credentials.access_key), - ("aws_secret_access_key", credentials.secret_key), - ("aws_session_token", credentials.token), - ) - if value is not None - } - if credentials is not None - else {} - ), + **{ # mutable-ok: merged into its mutable parent above + key: value + for key, value in ( + *credential_params, + ("aws_region_name", aws_region_name), + ) + if value is not None + }, } serves_via_rust: Final = rust_chat_completions_accepts( model=model, diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index 038ccf9b82f..e799ac7cf1f 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -329,29 +329,6 @@ class TestBedrockMoonshotBasic: # The model ID in the request body should be stripped assert transformed["model"] == "moonshot.kimi-k2-thinking" - def test_transform_request_bearer_token_skips_aws_credentials(self): - from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( - AmazonMoonshotConfig, - ) - - config = AmazonMoonshotConfig() - with patch.object( - config, - "get_credentials", - side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"), - ) as get_credentials: - transformed = config.transform_request( - model="bedrock/invoke/moonshot.kimi-k2-thinking", - messages=[{"role": "user", "content": "Hello"}], - optional_params={}, - litellm_params={"api_key": "bedrock-bearer-token"}, - headers={}, - ) - - get_credentials.assert_not_called() - assert transformed["model"] == "moonshot.kimi-k2-thinking" - - class TestBedrockMoonshotReasoningContent: """Tests for reasoning content extraction.""" 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..d3b1f5a6462 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py @@ -0,0 +1,20 @@ +from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, +) + + +def test_transform_request_bearer_token_skips_aws_credentials(): + 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", + "aws_region_name": "us-east-1", + }, + litellm_params={"api_key": "bedrock-bearer-token"}, + headers={}, + ) + + assert transformed["model"] == "moonshot.kimi-k2-thinking" 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 3a7ed9d806f..654aaa1afdf 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 @@ -138,14 +138,16 @@ def test_the_core_receives_the_credentials_this_handler_already_resolved(): def test_bearer_token_skips_aws_credential_resolution(): seen = _inject() - with patch.object( - BedrockConverseLLM, - "get_credentials", - side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"), - ) as get_credentials: - response = BedrockConverseLLM().completion(**_completion_kwargs(api_key="bedrock-bearer-token")) + response = BedrockConverseLLM().completion( + **_completion_kwargs( + api_key="bedrock-bearer-token", + optional_params={ + "maxTokens": 16, + "aws_profile_name": "litellm-profile-that-does-not-exist", + }, + ) + ) - get_credentials.assert_not_called() assert response.choices[0].message.content == "hello from rust" params = seen["call"][0]["optional_params"] assert params["aws_region_name"] == "us-east-1" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 74a55cc1ef2..f2baa958f01 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -1,5 +1,5 @@ import json -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, patch import pytest @@ -239,6 +239,34 @@ class TestBedrockAsyncInvokeEmbedding: assert status_response["status"] == "InProgress" assert "invocationArn" in status_response + @pytest.mark.asyncio + async def test_async_invoke_status_uses_sigv4_when_bearer_token_is_configured( + self, monkeypatch + ): + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") + response = Mock(status_code=200) + response.json.return_value = async_invoke_status_response + client = Mock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.llms.bedrock.embed.embedding.get_async_httpx_client", + return_value=client, + ): + status_response = await BedrockEmbedding()._get_async_invoke_status( + invocation_arn="arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456", + aws_region_name="us-east-1", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_session_token="test-session-token", + ) + + request_headers = client.get.await_args.kwargs["headers"] + assert request_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert status_response == async_invoke_status_response + def test_async_invoke_error_handling_missing_output_s3_uri(self): """Test error handling when output_s3_uri is missing for async invoke.""" from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import ( 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 114e473be98..33bbca1db1b 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -65,6 +65,7 @@ def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_re "client": client, "aws_region_name": "us-east-1", "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-east-1.amazonaws.com", + "aws_profile_name": "litellm-profile-that-does-not-exist", "api_key": test_api_key, } 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..6be43bed222 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 import httpx import pytest @@ -72,6 +73,24 @@ def test_get_config_class_stability_unchanged(): assert cls is BedrockStabilityImageEditConfig +def test_prepare_request_bearer_token_skips_aws_credentials(): + prepared_request = BedrockImageEdit()._prepare_request( + model="amazon.nova-canvas-v1:0", + image=[io.BytesIO(b"image")], + prompt="replace the background", + optional_params={ + "aws_profile_name": "litellm-profile-that-does-not-exist", + "aws_region_name": "us-east-1", + }, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + extra_headers=None, + logging_obj=Mock(), + api_key="bedrock-bearer-token", + ) + + assert prepared_request.prepped.headers["Authorization"] == "Bearer bedrock-bearer-token" + + def test_provider_config_router_returns_nova_for_canvas(): """ProviderConfigManager routes Nova Canvas to Nova image-edit config.""" cfg = get_bedrock_image_edit_config_for_model("amazon.nova-canvas-v1:0") diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 60b23defd1e..f824274e238 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -743,18 +743,14 @@ def test_bearer_token_skips_boto_credential_resolution( env_bearer_token: str | None, ): llm = BaseAWSLLM() - optional_params = {"aws_region_name": "us-east-1"} + optional_params = { + "aws_region_name": "us-east-1", + "aws_profile_name": "litellm-profile-that-does-not-exist", + } - with ( - patch( - "litellm.llms.bedrock.base_aws_llm.get_secret_str", - return_value=env_bearer_token, - ), - patch.object( - llm, - "get_credentials", - side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"), - ) as get_credentials, + with patch( + "litellm.llms.bedrock.base_aws_llm.get_secret_str", + return_value=env_bearer_token, ): credential_info = llm._get_boto_credentials_from_optional_params( optional_params, @@ -762,28 +758,30 @@ def test_bearer_token_skips_boto_credential_resolution( supports_bearer_token=True, ) - get_credentials.assert_not_called() assert credential_info.credentials is None assert credential_info.aws_region_name == "us-east-1" def test_shared_boto_helper_requires_credentials_by_default(): llm = BaseAWSLLM() - credentials = Credentials("test_key", "test_secret", "test_token") - with ( - patch( - "litellm.llms.bedrock.base_aws_llm.get_secret_str", - return_value="environment-bearer-token", - ), - patch.object(llm, "get_credentials", return_value=credentials) as get_credentials, + with patch( + "litellm.llms.bedrock.base_aws_llm.get_secret_str", + return_value="environment-bearer-token", ): credential_info = llm._get_boto_credentials_from_optional_params( - {"aws_region_name": "us-east-1"}, + { + "aws_region_name": "us-east-1", + "aws_access_key_id": "test_key", + "aws_secret_access_key": "test_secret", + "aws_session_token": "test_token", + }, ) - get_credentials.assert_called_once() - assert credential_info.credentials is credentials + assert credential_info.credentials is not None + assert credential_info.credentials.access_key == "test_key" + assert credential_info.credentials.secret_key == "test_secret" + assert credential_info.credentials.token == "test_token" def test_get_request_headers_with_env_var_bearer_token():