From 181a47b36f920c9e7432bfac88d53ac8f078dfad Mon Sep 17 00:00:00 2001 From: Kalai <101443484+likalight@users.noreply.github.com> Date: Mon, 24 Aug 2026 08:09:45 +0800 Subject: [PATCH 1/3] fix(bedrock): back-fill rerank document text when return_documents is set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bedrock's rerank API returns only `index` and `relevanceScore` and never echoes document text. LiteLLM exposes a Cohere-compatible interface where `return_documents` means the transformer back-fills `document.text` from the request's documents by index, but the Bedrock transformer ignored the parameter entirely — so `results[i]["document"]` was always absent and callers hit `KeyError: 'document'` no matter what they passed. Thread the `RerankRequest` into `_transform_response` (sync and async paths) and back-fill by index, the same way the HuggingFace rerank config already handles this identical provider limitation. Both `str` and `{"text": ...}` document shapes are supported. `request_data` is optional, so an existing caller passing only the response keeps its behaviour. Fixes #38006 Co-Authored-By: Claude Opus 5 (1M context) --- litellm/llms/bedrock/rerank/handler.py | 6 +- litellm/llms/bedrock/rerank/transformation.py | 39 +++++++++-- .../test_bedrock_rerank_return_documents.py | 70 +++++++++++++++++++ 3 files changed, 107 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 1cc72f265eb0..7fe068821d86 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -31,6 +31,7 @@ async def arerank( prepared_request: BedrockPreparedRequest, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, + request_data: RerankRequest | None = None, ): if client is None: client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) @@ -48,7 +49,7 @@ async def arerank( except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return BedrockRerankConfig()._transform_response(response.json()) + return BedrockRerankConfig()._transform_response(response.json(), request_data) def rerank( self, @@ -98,6 +99,7 @@ def rerank( if _is_async: return self.arerank( prepared_request, + request_data=request_data, timeout=timeout, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None, ) @@ -125,7 +127,7 @@ def rerank( response_json: Final = response.json() - return BedrockRerankConfig()._transform_response(response_json) + return BedrockRerankConfig()._transform_response(response_json, request_data) def _prepare_request( self, diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index 930dfa4dfbd7..97b44aa8d6d0 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/litellm/llms/bedrock/rerank/transformation.py @@ -22,6 +22,7 @@ RerankBilledUnits, RerankRequest, RerankResponse, + RerankResponseDocument, RerankResponseMeta, RerankResponseResult, RerankTokens, @@ -77,12 +78,29 @@ def _transform_request(self, request_data: RerankRequest) -> BedrockRerankReques sources=_sources, ) - def _transform_response(self, response: dict) -> RerankResponse: + @staticmethod + def _document_text(document: str | dict) -> str | None: + """Text of an input document, in either shape the API accepts.""" + if isinstance(document, str): + return document + if isinstance(document, dict): + text = document.get("text") + if isinstance(text, str): + return text + return None + + def _transform_response(self, response: dict, request_data: RerankRequest | None = None) -> RerankResponse: """ Transform the response from Bedrock into the RerankResponse format. example input: {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]} + + Bedrock never echoes document text, so `document` is back-filled from + the request's documents by index — the same thing the HuggingFace + rerank config does for the identical provider limitation. Without + `request_data` there is nothing to back-fill from, so the response is + left as-is. """ _billed_units = RerankBilledUnits(**response.get("usage", {"search_units": 1})) # by default 1 search unit _tokens: Final = RerankTokens(**response.get("usage", {})) @@ -92,13 +110,22 @@ def _transform_response(self, response: dict) -> RerankResponse: bedrock_results: Final = response.get("results") if bedrock_results: - _results = [ - RerankResponseResult( - index=result.get("index"), + # Cohere-compatible default: back-fill unless the caller opted out. + should_return_documents: Final = request_data is not None and request_data.return_documents is not False + original_documents: Final = request_data.documents if request_data is not None else [] + + _results = [] + for result in bedrock_results: + index = result.get("index") + _result = RerankResponseResult( + index=index, relevance_score=result.get("relevanceScore"), ) - for result in bedrock_results - ] + if should_return_documents and isinstance(index, int) and 0 <= index < len(original_documents): + text = self._document_text(original_documents[index]) + if text is not None: + _result["document"] = RerankResponseDocument(text=text) + _results.append(_result) if _results is None: raise ValueError(f"No results found in the response={response}") diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py new file mode 100644 index 000000000000..131df76e3e92 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py @@ -0,0 +1,70 @@ +"""Bedrock never echoes document text, so the transformer must back-fill it.""" + +import pytest + +from litellm.llms.bedrock.rerank.transformation import BedrockRerankConfig +from litellm.types.rerank import RerankRequest + +BEDROCK_RESPONSE = { + "results": [ + {"index": 1, "relevanceScore": 0.95}, + {"index": 0, "relevanceScore": 0.42}, + ] +} + +DOCUMENTS = ["the first passage", "the second passage"] + + +def _request(documents=None, return_documents=None): + return RerankRequest( + model="bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0", + query="which passage?", + documents=DOCUMENTS if documents is None else documents, + return_documents=return_documents, + ) + + +class TestBedrockRerankReturnDocuments: + def setup_method(self): + self.config = BedrockRerankConfig() + + @pytest.mark.parametrize("return_documents", [True, None]) + def test_back_fills_document_text_by_index(self, return_documents): + """`None` is the Cohere-compatible default and back-fills like `True`.""" + response = self.config._transform_response(BEDROCK_RESPONSE, _request(return_documents=return_documents)) + + # Ranked order is preserved, and each result carries its source text. + assert response.results[0]["index"] == 1 + assert response.results[0]["document"] == {"text": "the second passage"} + assert response.results[1]["index"] == 0 + assert response.results[1]["document"] == {"text": "the first passage"} + + def test_omits_documents_when_caller_opts_out(self): + response = self.config._transform_response(BEDROCK_RESPONSE, _request(return_documents=False)) + + assert all("document" not in result for result in response.results) + + def test_accepts_dict_documents(self): + response = self.config._transform_response( + BEDROCK_RESPONSE, + _request(documents=[{"text": "first as dict"}, {"text": "second as dict"}]), + ) + + assert response.results[0]["document"] == {"text": "second as dict"} + + def test_no_request_data_leaves_response_untouched(self): + """Older callers pass only the response; nothing to back-fill from.""" + response = self.config._transform_response(BEDROCK_RESPONSE) + + assert all("document" not in result for result in response.results) + + def test_out_of_range_index_is_skipped_not_raised(self): + response = self.config._transform_response({"results": [{"index": 7, "relevanceScore": 0.1}]}, _request()) + + assert response.results[0]["index"] == 7 + assert "document" not in response.results[0] + + def test_relevance_scores_are_preserved(self): + response = self.config._transform_response(BEDROCK_RESPONSE, _request()) + + assert [r["relevance_score"] for r in response.results] == [0.95, 0.42] From 3bcd861bc4a48a7720b0b469f58bd17c79282c5f Mon Sep 17 00:00:00 2001 From: Kalai <101443484+likalight@users.noreply.github.com> Date: Mon, 24 Aug 2026 08:19:33 +0800 Subject: [PATCH 2/3] refactor(bedrock): build rerank results functionally to satisfy the type gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first cut accumulated results with `_results = []` plus `.append`, and defaulted the documents list to a `[]` literal. Both are mutable-collection constructions (LIT002), so the type-discipline gate counted two new violations against the budget. Build the list in one shot from a comprehension over a per-result helper, and let `None` — rather than an empty list — mean "nothing to back-fill from", which folds the opt-out and the no-request case into one check. Annotate the new helpers with read-only views (Sequence/Mapping) rather than list/dict so they do not add LIT001 either. Behaviour is unchanged; the tests are untouched and still pass. transformation.py now sits below its base count on every gated rule (LIT001 5<6, LIT002 5<6, LIT010 3<4) and handler.py is unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- litellm/llms/bedrock/rerank/transformation.py | 55 +++++++++++++------ 1 file changed, 37 insertions(+), 18 deletions(-) diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index 97b44aa8d6d0..854e22f050dc 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/litellm/llms/bedrock/rerank/transformation.py @@ -4,6 +4,7 @@ Why separate file? Make it easy to see how transformation works """ +from collections.abc import Mapping, Sequence from typing import Final from litellm._uuid import uuid @@ -79,16 +80,40 @@ def _transform_request(self, request_data: RerankRequest) -> BedrockRerankReques ) @staticmethod - def _document_text(document: str | dict) -> str | None: - """Text of an input document, in either shape the API accepts.""" + def _document_text(documents: Sequence[str | Mapping[str, object]] | None, index: object) -> str | None: + """ + Text of the input document at `index`, in either shape the API accepts. + + Returns None when there is nothing to back-fill from: no documents, a + non-integer or out-of-range index, or a document carrying no text. + """ + if documents is None or not isinstance(index, int) or not 0 <= index < len(documents): + return None + document: Final = documents[index] if isinstance(document, str): return document if isinstance(document, dict): - text = document.get("text") + text: Final = document.get("text") if isinstance(text, str): return text return None + @classmethod + def _transform_result( + cls, result: Mapping[str, object], documents: Sequence[str | Mapping[str, object]] | None + ) -> RerankResponseResult: + """One Bedrock result, with `document` back-filled when available.""" + index: Final = result.get("index") + relevance_score: Final = result.get("relevanceScore") + text: Final = cls._document_text(documents, index) + if text is None: + return RerankResponseResult(index=index, relevance_score=relevance_score) + return RerankResponseResult( + index=index, + relevance_score=relevance_score, + document=RerankResponseDocument(text=text), + ) + def _transform_response(self, response: dict, request_data: RerankRequest | None = None) -> RerankResponse: """ Transform the response from Bedrock into the RerankResponse format. @@ -111,21 +136,15 @@ def _transform_response(self, response: dict, request_data: RerankRequest | None bedrock_results: Final = response.get("results") if bedrock_results: # Cohere-compatible default: back-fill unless the caller opted out. - should_return_documents: Final = request_data is not None and request_data.return_documents is not False - original_documents: Final = request_data.documents if request_data is not None else [] - - _results = [] - for result in bedrock_results: - index = result.get("index") - _result = RerankResponseResult( - index=index, - relevance_score=result.get("relevanceScore"), - ) - if should_return_documents and isinstance(index, int) and 0 <= index < len(original_documents): - text = self._document_text(original_documents[index]) - if text is not None: - _result["document"] = RerankResponseDocument(text=text) - _results.append(_result) + # None documents means "nothing to back-fill from", which covers + # both an opt-out and a caller that passed no request at all. + documents: Final = ( + request_data.documents + if request_data is not None and request_data.return_documents is not False + else None + ) + + _results = [self._transform_result(result, documents) for result in bedrock_results] if _results is None: raise ValueError(f"No results found in the response={response}") From 4c9d7f65b4d8a4dde26e999d8e987e0690fbd12a Mon Sep 17 00:00:00 2001 From: Kalai <101443484+likalight@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:32:38 +0800 Subject: [PATCH 3/3] test(bedrock): pin the rerank back-fill boundaries raised in review Two cases @kimnamu probed on #38006 while comparing this PR with #38007, now pinned so neither can regress silently: - A dict document with no `text` key reaches Bedrock as a `jsonDocument` and `RerankResponseDocument` carries only `text`, so there is nowhere to put it. No back-fill, but no crash: the caller still gets index and score. This is the one case where the reported `KeyError: 'document'` survives, so it is a documented boundary rather than an oversight. - A non-integer `index` is tolerated rather than raising, matching main's behaviour (pydantic coerces "2" to 2 upstream). Co-Authored-By: Claude Opus 5 (1M context) --- .../test_bedrock_rerank_return_documents.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py index 131df76e3e92..f085487bf543 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py @@ -64,6 +64,27 @@ def test_out_of_range_index_is_skipped_not_raised(self): assert response.results[0]["index"] == 7 assert "document" not in response.results[0] + def test_json_documents_yield_no_document_and_do_not_raise(self): + """ + A dict document with no `text` key goes to Bedrock as a `jsonDocument` + (`_transform_sources`), and `RerankResponseDocument` carries only + `text`, so there is nowhere to put it. Pinned as a known boundary: no + back-fill, but no crash either. Raised by @kimnamu on #38006. + """ + response = self.config._transform_response( + BEDROCK_RESPONSE, + _request(documents=[{"title": "zero", "body": "b0"}, {"title": "one", "body": "b1"}]), + ) + + assert all("document" not in result for result in response.results) + assert [r["index"] for r in response.results] == [1, 0] + + def test_non_integer_index_is_tolerated(self): + """`index` is coerced by pydantic upstream; do not raise on it here.""" + response = self.config._transform_response({"results": [{"index": "2", "relevanceScore": 0.9}]}, _request()) + + assert response.results[0]["relevance_score"] == 0.9 + def test_relevance_scores_are_preserved(self): response = self.config._transform_response(BEDROCK_RESPONSE, _request())