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..90378d33d3a 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 ### @@ -369,14 +373,21 @@ 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, **{ # 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), + *credential_params, ("aws_region_name", aws_region_name), ) if value is not None 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..e799ac7cf1f 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -329,7 +329,6 @@ class TestBedrockMoonshotBasic: # The model ID in the request body should be stripped 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 8e67a7e3438..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 @@ -135,6 +135,27 @@ 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() + + response = BedrockConverseLLM().completion( + **_completion_kwargs( + api_key="bedrock-bearer-token", + optional_params={ + "maxTokens": 16, + "aws_profile_name": "litellm-profile-that-does-not-exist", + }, + ) + ) + + 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/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 50e2b53c2b3..f824274e238 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,63 @@ 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", + "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, + ): + credential_info = llm._get_boto_credentials_from_optional_params( + optional_params, + api_key=api_key, + supports_bearer_token=True, + ) + + 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() + + 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_access_key_id": "test_key", + "aws_secret_access_key": "test_secret", + "aws_session_token": "test_token", + }, + ) + + 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(): # Setup llm = BaseAWSLLM() - credentials = Credentials("test_key", "test_secret", "test_token") + credentials = None headers = {"Content-Type": "application/json"} headers_dict = headers.copy() @@ -831,7 +884,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"