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():