From 33faceb455d8f314391b4e69a8870273b6c0c4cb Mon Sep 17 00:00:00 2001 From: Rabari Mihir <169349389+Mihir-Rabari@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:35:04 +0000 Subject: [PATCH] feat(bedrock): authenticate rerank with a Bedrock API key as a bearer token The rerank handler hardcoded SigV4 and ignored the api_key set on the deployment, so a user with only a Bedrock API key fell back to the default AWS credential chain and failed with a bedrock:Rerank permission error. The Cohere Rerank 3.5 model card documents calling the Rerank API with a long-term API key, and chat, embeddings and image generation already send it as Authorization: Bearer, so rerank now does the same and keeps SigV4 only when no key is configured --- litellm/llms/bedrock/rerank/handler.py | 11 +++- litellm/rerank_api/main.py | 1 + .../test_bedrock_rerank_header_forwarding.py | 57 +++++++++++++++++-- 3 files changed, 60 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 01ef5af2b67..ef282f1e2db 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -16,7 +16,7 @@ from litellm.types.llms.bedrock import BedrockPreparedRequest from litellm.types.rerank import RerankRequest from litellm.types.utils import RerankResponse -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token from ..common_utils import BedrockError from .transformation import BedrockRerankConfig @@ -76,6 +76,7 @@ class BedrockRerankHandler(BaseAWSLLM): api_base: str | None = None, extra_headers: dict | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, + api_key: str | None = None, ) -> RerankResponse: request_data: Final = RerankRequest( model=model, @@ -93,6 +94,7 @@ class BedrockRerankHandler(BaseAWSLLM): api_base=api_base, extra_headers=extra_headers, data=cast(dict, data), + api_key=api_key, ) logging_obj.pre_call( @@ -150,8 +152,11 @@ class BedrockRerankHandler(BaseAWSLLM): extra_headers: dict | None, data: dict, optional_params: dict, + api_key: str | None = None, ) -> BedrockPreparedRequest: - 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, bearer_token=bedrock_bearer_token(api_key) + ) ### SET RUNTIME ENDPOINT ### _, proxy_endpoint_url = self.get_runtime_endpoint( @@ -174,7 +179,7 @@ class BedrockRerankHandler(BaseAWSLLM): endpoint_url=proxy_endpoint_url, data=body, headers=headers, - supports_bearer_token=False, + api_key=api_key, ) return BedrockPreparedRequest( diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index f7e2616324f..c747131bb8f 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -381,6 +381,7 @@ def rerank( optional_params=optional_params.model_dump(exclude_unset=True), timeout=optional_params.timeout, api_base=api_base, + api_key=dynamic_api_key or optional_params.api_key, extra_headers=merged_headers, logging_obj=litellm_logging_obj, client=client, diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 495e0f56c84..67a25335bde 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -451,11 +451,12 @@ def test_bedrock_rerank_forwarded_headers_excluded_from_sigv4_signature(): 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): +def test_bedrock_rerank_uses_bedrock_api_key_as_bearer(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. + The Cohere Rerank 3.5 model card documents calling the Rerank API with a + long-term Bedrock API key (docs.aws.amazon.com model-card-cohere-rerank-3-5, + checked 2026-10-16), so a key in AWS_BEARER_TOKEN_BEDROCK must be sent as + an Authorization: Bearer header and must not trigger SigV4 signing """ monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "test-bedrock-api-key") @@ -476,11 +477,55 @@ def test_bedrock_rerank_signs_with_sigv4_even_when_bedrock_api_key_is_set(monkey 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]}" + assert authorization == "Bearer test-bedrock-api-key", ( + f"rerank must authenticate with the Bedrock API key, got Authorization={authorization[:30]}" ) +def test_bedrock_rerank_with_api_key_authenticates_with_bearer_not_sigv4(): + """ + A Bedrock API key configured on the deployment must authenticate the rerank + call with Authorization: Bearer, matching how Bedrock chat and embeddings + already treat it. SigV4 signing with the instance role ignores the key + """ + client = HTTPHandler() + model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0" + + with ( + patch.object(client, "post") as mock_post, + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport + "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", + return_value=create_mock_credentials(), + ), + patch("botocore.auth.SigV4Auth") as mock_sigv4, + ): + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(bedrock_rerank_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_response.raise_for_status = lambda: None + mock_post.return_value = mock_response + + response = litellm.rerank( + model=model, + query=test_query, + documents=test_documents, + top_n=3, + client=client, + api_key="test-bearer-key-12345", + aws_region_name="us-east-1", + ) + + assert isinstance(response, litellm.RerankResponse) + + sent_headers = mock_post.call_args.kwargs["headers"] + assert sent_headers.get("Authorization") == "Bearer test-bearer-key-12345", ( + "rerank with a Bedrock API key must send Authorization: Bearer, " + f"got headers={sent_headers}" + ) + assert not mock_sigv4.called, "rerank with a Bedrock API key must not SigV4-sign" + + @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