Merge 4de35621dd into 39e5b0c2d1

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:

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

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)

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(

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)

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)

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."""
 

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"

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

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 (

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,
         }
 

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")

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"