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

[likalight]

Fixes #38006.

The bug

Bedrock's rerank API returns only index and relevanceScore — it never echoes
document text. LiteLLM exposes a Cohere-compatible interface where
return_documents means the response transformer back-fills document.text
from the request's documents by index.

BedrockRerankConfig._transform_response never saw the request at all:

def _transform_response(self, response: dict) -> RerankResponse:

so it could not back-fill even in principle, and return_documents was a
no-op. results[i]["document"] was always absent, and a RAG pipeline using
Bedrock rerank had to re-join ranked indices to source text by hand.

The fix

Thread the RerankRequest through to _transform_response (both the sync path
and arerank) and back-fill by index — the same approach
HuggingFaceRerankConfig already uses for this identical provider limitation.

• Both document shapes are handled: plain str and {"text": ...}.
• request_data is optional, so any existing caller passing only the
  response keeps today's behaviour rather than breaking.
• An out-of-range index is skipped rather than raising, so a malformed response
  degrades instead of crashing the call.

The None default — settled in review

RerankRequest.return_documents defaults to None, not True, so I originally
raised this as a judgement call. @kimnamu measured it on the issue and the
answer is clear: rerank(...) declares return_documents: bool | None = True
while arerank(...) declares None and forwards it positionally, so the
transform layer sees True on the sync path and None on the async path.
Treating None as back-fill is what keeps the two agreeing; treating it as
opt-out splits them, and closing that split from the other side would change
what the async path sends to other providers (fireworks_ai forwards the flag
only when it is not None).

Note this is a deliberate divergence from Cohere, whose v1 rerank defaults
return_documents to false, and from the HuggingFace config here, which treats
absent as False. Bedrock ends up more generous than both. Say the word if you
would rather match them and accept the sync/async split.

Release note worth making

The proxy's /v1/rerank runs through route_type="arerank", so proxy callers
who never set the flag are on the None branch and will start receiving
document text where they previously got none. For 50 documents of ~1 KB
passages that is a serialized response of 2,217 → 54,107 bytes. That is the fix
working as intended, but it is a visible payload change worth calling out in
release notes.

Known boundary: JSON documents

A dict document with no text key is sent to Bedrock as a jsonDocument
(_transform_sources), and RerankResponseDocument carries only text, so
there is nowhere to put it. Those callers still see no document key. There is
a test pinning this explicitly — no back-fill, but no crash, so index and score
still come back. Serializing the JSON into text would close it, but that
invents a shape neither Cohere nor the HuggingFace config produces, so I have
left it as a documented limitation rather than decide it inside a bug fix.

Verification

Seven tests in tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py:

• back-fills by index for return_documents=True and None, with ranked
  order preserved (the response ranks index 1 above index 0, so this also pins
  that text is matched by index rather than by position)
• omits document entirely when the caller passes False
• handles {"text": ...} document dicts
• leaves the response untouched when no request_data is supplied
• skips an out-of-range index instead of raising
• preserves relevance scores

6 of the 7 fail without the source change. The existing
test_bedrock_rerank_header_forwarding.py (which exercises arerank) still
passes — 14/14 in that directory — confirming the signature change is safe
on the async path.

ruff check and ruff format --check clean on the touched files.

Unrelated pre-existing failures

tests/test_litellm/llms/bedrock/ has 10 failures on this branch, but the same
10 fail on a clean checkout with nothing applied (verified by stashing) — mostly
in test_anthropic_claude3_transformation.py, plus
test_sigv4_matches_rust_golden_vector and test_mantle.py. Untouched by this
PR; flagging in case they are not already known.

---

On the red code-quality check: it is not from this PR. It fails on every
PR against litellm_internal_staging right now, including ones that touch no
workflow files, because three unit shards in .github/workflows/test-unit.yml
cap the job below the startup-safety invariant. Fixed independently in #38046;
this PR needs no change for it.

[CLAassistant]

All committers have signed the CLA.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR combines a Bedrock rerank response fix with broader CI, telemetry, spend-log session, and Responses API streaming changes

• Back-fills Bedrock rerank document text from the original request
• Adds tenant-routed New Relic OpenTelemetry support and callback configuration validation
• Reorganizes CI suites and coverage enforcement
• Changes Responses API stream identifiers, tool events, and session spend-log availability

Confidence Score: 4/5

The streaming regression should be fixed before merging because slow providers can withhold every initial Responses API event until their first token

Both iterator modes now consume the upstream stream before yielding response.created, exposing otherwise healthy long-time-to-first-token requests to client idle and first-byte timeouts

Files Needing Attention: litellm/responses/litellm_completion_transformation/streaming_iterator.py, tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py

Important Files Changed

Filename	Overview
litellm/responses/litellm_completion_transformation/streaming_iterator.py	Adds response-ID priming and expanded event assembly, but now waits for provider output before emitting initial stream events
litellm/llms/bedrock/rerank/transformation.py	Correctly restores requested source text by validated Bedrock result index while preserving opt-out behavior
litellm/llms/bedrock/rerank/handler.py	Threads the normalized rerank request through both synchronous and asynchronous response transformation
litellm/integrations/otel/plumbing/routing.py	Extends tenant tracer cache identity with fixed-region endpoints and confines dynamic credentials to the owning exporter
litellm/proxy/utils.py	Adds an event-driven spend-log flush wake-up used to improve Responses API session lookup availability
.github/workflows/test-postgres.yml	Ports proxy behavior, security, and migration suites from removed CircleCI jobs into a GitHub Actions matrix

Comments Outside Diff (1)

1. litellm/responses/litellm_completion_transformation/streaming_iterator.py, line 369-375 (link)

   Initial stream events are blocked

   When a provider has a long time to first token, _aprime_response_id waits for its first chunk before emitting response.created, causing client first-byte or idle timeouts.

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

[greptile-apps]

Regression test file is fragmented

Move this bug-fix coverage into the existing mapped transformation.py test file and remove redundant explanatory comments to keep maintenance localized.

Context Used: CLAUDE.md (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing likalight:fix/bedrock-rerank-return-documents (4c9d7f6) 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 for pinning the jsonDocument and non-integer index boundaries in 4c9d7f6, and for folding the sync/async measurement into the description. I re-ran the branch on top of today's litellm_internal_staging (02522a54) and the fix still does exactly what it should; the only thing between this PR and green is a one-hunk conflict that landed after you opened it.

What changed underneath: #35419 (bb227420, Aug 25) threads logging_obj through BedrockRerankHandler.arerank on the same lines this PR threads request_data, so GitHub now reports mergeable_state: dirty (same for #38007). Keeping both kwargs is the whole resolution; the arerank signature itself merges cleanly (logging_obj positional after prepared_request, request_data stays the trailing kwarg):

# litellm/llms/bedrock/rerank/handler.py, after merging litellm_internal_staging
return self.arerank(
    prepared_request,
    logging_obj=logging_obj,    # from #35419
    request_data=request_data,  # this PR
    timeout=timeout,
    client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None,
)

Tests on the merged tree, tests/test_litellm/llms/bedrock/rerank/:

Tree	Result
litellm_internal_staging @ 02522a54, no PR	10 passed
PR 4c9d7f6 merged into staging	19 passed
same, with litellm/llms/bedrock/rerank/ reverted to staging (PR tests kept)	8 failed, 11 passed
ruff check + ruff format --check on the touched files	clean

The one PR test that passes without the fix is test_no_request_data_leaves_response_untouched, which is the intended no-op path.

End to end, only HTTPHandler.post / AsyncHTTPHandler.post patched (real SigV4 signing with dummy keys, return_documents not passed):

=== litellm_internal_staging @ 02522a54 ===
rerank  [{"index": 1, "document": null}, {"index": 0, "document": null}]
arerank [{"index": 1, "document": null}, {"index": 0, "document": null}]
=== PR 4c9d7f65 merged into staging ===
rerank  [{"index": 1, "document": {"text": "beta passage"}}, {"index": 0, "document": {"text": "alpha passage"}}]
arerank [{"index": 1, "document": {"text": "beta passage"}}, {"index": 0, "document": {"text": "alpha passage"}}]

So the arerank path (what the proxy's /v1/rerank goes through) still back-fills after #35419's logging_obj plumbing, and the sync and async entry points agree with the flag unset, as measured on the issue.

Repro script

import asyncio, json, httpx, litellm
from unittest.mock import patch
from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler

BODY = {"results": [{"index": 1, "relevanceScore": 0.9}, {"index": 0, "relevanceScore": 0.1}]}
def fake(*a, **k):
    return httpx.Response(200, json=BODY, request=httpx.Request("POST", "https://bedrock-agent-runtime.us-east-1.amazonaws.com/rerank"))
async def afake(*a, **k):
    return fake()

kw = dict(model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", query="q",
          documents=["alpha passage", "beta passage"], aws_region_name="us-east-1",
          aws_access_key_id="AKIAIOSFODNN7EXAMPLE", aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY")
with patch.object(HTTPHandler, "post", side_effect=fake), patch.object(AsyncHTTPHandler, "post", side_effect=afake):
    s = litellm.rerank(**kw)
    a = asyncio.run(litellm.arerank(**kw))
for name, r in (("rerank ", s), ("arerank", a)):
    print(name, json.dumps([{"index": x["index"], "document": x.get("document")} for x in r.results]))

Nothing else from my side; once the branch is rebased or merged with staging this looks ready for a maintainer pass.