fix(bedrock/rerank): populate document text when return_documents is true

[AbhAy120204]

TLDR

Problem this solves:

• Bedrock rerank ignores return_documents — results never carry document.text
• Callers using the Cohere-compatible default (return_documents=True) get a KeyError on results[i]["document"]

How it solves it:

• Threads documents and return_documents into the response transformer
• Back-fills document.text from the original input list by index, same as the HuggingFace rerank config

User Flow

Before: a RAG pipeline using Bedrock rerank to retrieve ranked source passages gets no document text back

1. They call litellm.rerank(model="bedrock/arn:...", query="...", documents=["doc 0", "doc 1", ...], return_documents=True)
2. The response comes back with results: [{"index": 1, "relevance_score": 0.95}], no document field
3. They access results[0]["document"]["text"] and get KeyError: 'document'

After: the same call returns document text in each result

1. Same call
2. The response comes back with results: [{"index": 1, "relevance_score": 0.95, "document": {"text": "doc 1"}}]
3. results[0]["document"]["text"] returns the text of the highest-ranked source document

Relevant issues

Fixes #38006

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• The handful of test files covering my change pass locally, e.g. uv run pytest tests/test_litellm/llms/bedrock/rerank/ -v. Leave the suites (make test-unit-*, make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
• My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

This is a pure SDK transformation bug: Bedrock's API only returns {index, relevanceScore} and never echoes document text. The fix is client-side back-fill by index, confirmed by the regression test below. No live Bedrock call is needed to reproduce or verify — the transformer is fully exercisable without a provider.

Shared setup: uv run pytest tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py -v

Before (ff02d5c)

1. FAILED test_bedrock_rerank_return_documents_true_populates_document_text
AssertionError: return_documents=True must populate document, got {'index': 2, 'relevance_score': 0.95}
2. FAILED test_bedrock_rerank_return_documents_true_async
KeyError: 'document'
3. 2 failed, 1 passed

After (8fb4545)

1. PASSED test_bedrock_rerank_return_documents_false_omits_document_field
2. PASSED test_bedrock_rerank_return_documents_true_async
3. PASSED test_bedrock_rerank_return_documents_true_populates_document_text
4. 3 passed in 1.08s

Type

🐛 Bug Fix

Caveats (if any)

• Bedrock's API never returns document text in its response, so back-fill is always from the original request documents list. This matches the existing HuggingFace rerank config which handles the same provider limitation the same way (litellm/llms/huggingface/rerank/transformation.py:234-247)
• Response size grows when back-fill kicks in. Proxy /v1/rerank routes through arerank where return_documents defaults to None → back-fill on, so callers who never set the flag will start getting document text. For 50 documents of ~1 KB passages the serialized response goes from ~2 KB to ~54 KB. Worth a release note
• Documents passed as dicts without a text key are sent to Bedrock as jsonDocument, and RerankResponseDocument carries only text. Those results have no back-fill target and come back with only index and relevance_score (no crash). The HuggingFace rerank config has the same limitation

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[greptile-apps]

Greptile Summary

The PR propagates the original rerank documents and return_documents setting through Bedrock’s synchronous and asynchronous response paths.

• Backfills each valid string document into the normalized result using Bedrock’s returned index.
• Omits document data when return_documents is false.
• Adds synchronous and asynchronous regression coverage for the new behavior.

Confidence Score: 4/5

The PR appears safe to merge, with only non-blocking repository-convention cleanup needed in the transformer and tests.

The synchronous and asynchronous Bedrock paths consistently propagate the new inputs and correctly backfill valid string documents; the remaining concern is avoidable mutation and class-level test patching.

Files Needing Attention: litellm/llms/bedrock/rerank/transformation.py; tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py

Important Files Changed

Filename	Overview
litellm/llms/bedrock/rerank/handler.py	Threads documents and the effective return-documents setting through both Bedrock response paths without leaving an inconsistent in-repository caller.
litellm/llms/bedrock/rerank/transformation.py	Backfills string documents safely by validated result index, but uses mutable construction contrary to repository conventions.
tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py	Covers enabled, disabled, and asynchronous behavior, while relying on class-level monkeypatching rather than dependency injection.

_{Reviews (1): Last reviewed commit: "fix(bedrock/rerank): populate document t..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[AbhAy120204]

The code-quality failure is on the pre-existing check_workflow_startup_safety step which is failing on litellm_internal_staging itself (the .github/workflows/test-unit.yml file has an unrelated job-timeout vs pytest-timeout mismatch introduced by upstream commits). This PR touches no workflow files.

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing AbhAy120204:litellm_fix_bedrock_rerank_return_documents (8fb4545) with litellm_internal_staging (d447be1)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (f005afa) during the generation of this report, so d447be1 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[kimnamu]

Thanks @AbhAy120204 for turning both of those around the same morning, and for pinning them with tests rather than only fixing them. I re-ran 8fb4545d and everything you listed holds:

Check on 8fb4545d	Result
inspect.signature(BedrockRerankHandler.arerank)	(prepared_request, timeout=None, client=None, documents=None, return_documents=True)
arerank(prepared_request, 10.0, client) positional, post mocked	returns results, no error
response {"index": "2", "relevanceScore": 0.9}	[{'index': 2, 'relevance_score': 0.9}], no raise
tests/test_litellm/llms/bedrock/rerank/	14 passed
ruff check (0.15.3) and ruff format --check on the three touched files	clean

One thing landed underneath since then. #35419 (bb227420, Aug 25) added logging_obj as a required positional parameter to arerank and passes it in the same self.arerank(...) call this PR extends, so GitHub now reports the branch CONFLICTING against litellm_internal_staging. It is one hunk:

<<<<<<< HEAD
                documents=documents,
                return_documents=should_return_documents,
=======
                logging_obj=logging_obj,
>>>>>>> litellm_internal_staging

Keeping all three kwargs is the whole resolution. On that merged tree (staging 168a0055 + this PR):

Tree	tests/test_litellm/llms/bedrock/rerank/
staging alone	10 passed
staging + this PR, hunk resolved as above	17 passed
same, with litellm/llms/bedrock/rerank/ reverted to staging (PR tests kept)	5 failed, 12 passed

The end-to-end shape is unchanged after the merge: rerank and arerank with the flag unset or True carry document.text, False omits it, on both paths.

One consequence of #35419 before you resolve: the pre-fix positional arity is already gone upstream. After the merge the signature is (prepared_request, logging_obj, timeout=None, client=None, documents=None, return_documents=True), so arerank(prepared_request, 10.0, client) fails inside client.post because 10.0 lands in logging_obj. That is upstream's change, not this PR's, so the trailing-kwargs placement you chose still looks right; there is just no old arity left to protect.

Merged-tree probe output

1) arerank signature: (self, prepared_request: BedrockPreparedRequest, logging_obj: Logging, timeout: float | Timeout | None = None, client: AsyncHTTPHandler | None = None, documents: Sequence[str | dict[str, Any]] | None = None, return_documents: bool = True)
2) non-int index: [{'index': 2, 'relevance_score': 0.9}]
3) positional arerank(prepared, timeout, client): FAILS -> ValidationError (10.0 lands in logging_obj, client in timeout)
4) e2e results[0]:
   rerank  return_documents=unset -> {'index': 2, 'relevance_score': 0.95, 'document': {'text': 'doc two text'}}
   arerank return_documents=unset -> {'index': 2, 'relevance_score': 0.95, 'document': {'text': 'doc two text'}}
   rerank  return_documents=True  -> {'index': 2, 'relevance_score': 0.95, 'document': {'text': 'doc two text'}}
   arerank return_documents=True  -> {'index': 2, 'relevance_score': 0.95, 'document': {'text': 'doc two text'}}
   rerank  return_documents=False -> {'index': 2, 'relevance_score': 0.95}
   arerank return_documents=False -> {'index': 2, 'relevance_score': 0.95}