fix(bedrock): back-fill rerank document text when return_documents is set

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,

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
@@ -22,6 +23,7 @@
     RerankBilledUnits,
     RerankRequest,
     RerankResponse,
+    RerankResponseDocument,
     RerankResponseMeta,
     RerankResponseResult,
     RerankTokens,
@@ -77,12 +79,53 @@ def _transform_request(self, request_data: RerankRequest) -> BedrockRerankReques
             sources=_sources,
         )
 
-    def _transform_response(self, response: dict) -> RerankResponse:
+    @staticmethod
+    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: 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.
 
         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 +135,16 @@ def _transform_response(self, response: dict) -> RerankResponse:
 
         bedrock_results: Final = response.get("results")
         if bedrock_results:
-            _results = [
-                RerankResponseResult(
-                    index=result.get("index"),
-                    relevance_score=result.get("relevanceScore"),
-                )
-                for result in bedrock_results
-            ]
+            # Cohere-compatible default: back-fill unless the caller opted out.
+            # 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}")

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

@@ -0,0 +1,91 @@
+"""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_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())
+
+        assert [r["relevance_score"] for r in response.results] == [0.95, 0.42]