Merge e47e989341 into cdb60af024

litellm/llms/bedrock/base_aws_llm.py

@@ -1434,9 +1434,12 @@ class BaseAWSLLM:
         data: str | bytes,
         headers: dict,
         api_key: str | None = None,
+        supports_bearer_token: bool = True,
     ) -> AWSPreparedRequest:
-        if api_key is not None:
-            aws_bearer_token: str | None = api_key
+        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")
 

litellm/llms/bedrock/rerank/handler.py

@@ -138,11 +138,6 @@ class BedrockRerankHandler(BaseAWSLLM):
         data: dict,
         optional_params: dict,
     ) -> BedrockPreparedRequest:
-        try:
-            from botocore.auth import SigV4Auth
-            from botocore.awsrequest import AWSRequest
-        except ImportError:
-            raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
         boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)
 
         ### SET RUNTIME ENDPOINT ###
@@ -153,24 +148,21 @@ class BedrockRerankHandler(BaseAWSLLM):
         )
         proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime")
         proxy_endpoint_url = f"{proxy_endpoint_url}/rerank"
-        sigv4: Final = SigV4Auth(
-            boto3_credentials_info.credentials,
-            "bedrock",
-            boto3_credentials_info.aws_region_name,
-        )
-        # Make POST Request
-        body: Final = json.dumps(data).encode("utf-8")
 
+        body: Final = json.dumps(data).encode("utf-8")
         headers = {"Content-Type": "application/json"}
         if extra_headers is not None:
             headers = {"Content-Type": "application/json", **extra_headers}
-        request: Final = AWSRequest(method="POST", url=proxy_endpoint_url, data=body, headers=headers)
-        sigv4.add_auth(request)
-        if (
-            extra_headers is not None and "Authorization" in extra_headers
-        ):  # prevent sigv4 from overwriting the auth header
-            request.headers["Authorization"] = extra_headers["Authorization"]
-        prepped: Final = request.prepare()
+
+        prepped: Final = self.get_request_headers(
+            credentials=boto3_credentials_info.credentials,
+            aws_region_name=boto3_credentials_info.aws_region_name,
+            extra_headers=extra_headers,
+            endpoint_url=proxy_endpoint_url,
+            data=body,
+            headers=headers,
+            supports_bearer_token=False,
+        )
 
         return BedrockPreparedRequest(
             endpoint_url=proxy_endpoint_url,

tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py

@@ -12,6 +12,7 @@ import pytest
 
 import litellm
 from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo
+from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 
 # Mock response for Bedrock rerank
@@ -402,6 +403,66 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
             pytest.fail(f"Failed to merge and forward headers: {str(e)}")
 
 
+def test_bedrock_rerank_forwarded_headers_excluded_from_sigv4_signature():
+    """
+    A forwarded header like x-forwarded-for can be rewritten between LiteLLM
+    signing the request and AWS receiving it (e.g. by an intermediate load
+    balancer), which invalidates the signature if that header was part of
+    the signed set. It must still reach Bedrock, just unsigned.
+    """
+    handler = BedrockRerankHandler()
+
+    prepared_request = handler._prepare_request(
+        model="cohere.rerank-v3-5:0",
+        api_base=None,
+        extra_headers={"x-forwarded-for": "203.0.113.5"},
+        data={"query": test_query, "documents": test_documents},
+        optional_params={
+            "aws_access_key_id": "test-access-key",
+            "aws_secret_access_key": "test-secret-key",
+            "aws_region_name": "us-east-1",
+        },
+    )
+
+    headers = prepared_request["prepped"].headers
+    signed_headers = headers["Authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")
+
+    assert "x-forwarded-for" not in signed_headers, (
+        f"x-forwarded-for must not be part of the SigV4 signature, got SignedHeaders={signed_headers}"
+    )
+    assert headers["x-forwarded-for"] == "203.0.113.5", "forwarded header must still reach Bedrock, unsigned"
+
+
+def test_bedrock_rerank_signs_with_sigv4_even_when_bedrock_api_key_is_set(monkeypatch):
+    """
+    Bedrock API keys are only valid for Bedrock and Bedrock Runtime actions, not for
+    Agents for Amazon Bedrock Runtime ones. Rerank is served by bedrock-agent-runtime,
+    so it has to keep signing with SigV4 even when AWS_BEARER_TOKEN_BEDROCK is set.
+    """
+    monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "test-bedrock-api-key")
+
+    handler = BedrockRerankHandler()
+
+    prepared_request = handler._prepare_request(
+        model="cohere.rerank-v3-5:0",
+        api_base=None,
+        extra_headers=None,
+        data={"query": test_query, "documents": test_documents},
+        optional_params={
+            "aws_access_key_id": "test-access-key",
+            "aws_secret_access_key": "test-secret-key",
+            "aws_region_name": "us-east-1",
+        },
+    )
+
+    assert prepared_request["endpoint_url"].startswith("https://bedrock-agent-runtime.")
+
+    authorization = prepared_request["prepped"].headers["Authorization"]
+    assert authorization.startswith("AWS4-HMAC-SHA256"), (
+        f"rerank must sign with SigV4, got Authorization={authorization[:30]}"
+    )
+
+
 @pytest.mark.asyncio
 async def test_bedrock_rerank_records_llm_api_duration():
     """The bedrock rerank handler must feed httpx timing into the logging obj, so the