fix(bedrock): support tool search results + chat annotations by Sameerlite · Pull Request #29120 · BerriAI/litellm

Sameerlite · 2026-05-28T07:36:55Z

Summary

add optional search_results support on chat-completions tool messages and map it to Bedrock toolResult.content[].searchResult blocks
make Bedrock converse response transformation expose citations as OpenAI-style message.annotations while keeping provider-specific citation payloads
add regression tests for both tool-result search result conversion and citations-to-annotations response mapping
Fixes LIT-3401

Test plan

pytest -q tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py -k \"tool_message_search_results_maps_to_bedrock_search_result_block\"
pytest -q tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py -k \"search_results_maps_to_bedrock_search_result_block or citations_content_maps_to_annotations\"
python -m black --check litellm/llms/bedrock/chat/converse_transformation.py litellm/types/llms/openai.py litellm/types/llms/bedrock.py litellm/litellm_core_utils/prompt_templates/factory.py tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py enterprise/litellm_enterprise/proxy/hooks/managed_files.py

Note

Medium Risk
Touches request/response translation on the Bedrock Converse path and proxy passthrough headers; behavior changes for tool messages with search_results and citation-heavy assistant output, with regression tests but some risk of edge-case mismatches with provider payloads.

Overview
Adds Bedrock Converse ↔ OpenAI chat bridging for RAG-style tool output and grounded replies.

Tool messages: Optional search_results on chat-completions tool messages maps to Bedrock toolResult.content[].searchResult (and sets status: success), preferring that path over normal content so mixed text/search blocks are avoided. Types gain SearchResultBlock / searchResult on tool result content.

Assistant responses: Nova/Bedrock citationsContent is converted to OpenAI-style message.annotations (url_citation with span indices, title, URL), with citation text merged into content when the model only returns punctuation in text blocks. ChatCompletionResponseMessage is typed for annotations.

Proxy passthrough: Custom response headers are filtered with the same exclusion list as upstream headers so defaults like content-length: 0 cannot override real passthrough bodies.

Regression tests cover search-result mapping, citations→annotations, passthrough headers, and managed-file ID scanning when user content includes Bedrock toolResult blocks without OpenAI type fields (large formatting-only churn in enterprise managed-files tests).

^{Reviewed by Cursor Bugbot for commit 65f8e21. Bugbot is set up for automated code reviews on this repo. Configure here.}

…notations Add an optional tool-message search_results path that maps directly to Bedrock toolResult.searchResult blocks, and convert Converse citationsContent into chat completion annotations for user-facing citation metadata. Co-authored-by: Cursor <cursoragent@cursor.com>

codecov · 2026-05-28T07:40:46Z

Codecov Report

❌ Patch coverage is 93.37748% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...llm/litellm_core_utils/prompt_templates/factory.py	92.30%	4 Missing ⚠️
...tellm/llms/bedrock/chat/converse_transformation.py	92.15%	4 Missing ⚠️
...itellm_core_utils/prompt_templates/common_utils.py	91.66%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

greptile-apps · 2026-05-28T07:41:16Z

Greptile Summary

This PR bridges Bedrock Converse ↔ OpenAI chat for two RAG-oriented features: tool messages can now carry a search_results extension mapped to Bedrock toolResult.content[].searchResult blocks, and Nova citationsContent responses are converted to OpenAI-style message.annotations.

Tool message search_results: _build_bedrock_tool_result_content_blocks prefers search_results over content when the field is present and non-empty, sets status: success, and falls back correctly to normal content processing for empty or absent lists. Empty-list fallback is covered by a new regression test.
Citations → annotations: _transform_citations_to_annotations maps citationsContent blocks to url_citation annotations with span offsets tracked across multiple blocks; annotations are only attached when citations text is stitched into content_str (empty or punctuation-only content case), with a test verifying they are omitted for mixed alphanumeric + citations responses.
Pass-through headers: Custom response headers are now filtered through the same excluded_headers set as upstream headers, preventing framework defaults like content-length: 0 from overriding real passthrough bodies.

Confidence Score: 5/5

Safe to merge; changes are additive, the only behavioral delta on existing callers is the pass-through header filter (conservative and beneficial), and the new Bedrock paths are gated behind an optional search_results field.

All three code paths have regression tests, the empty-list fallback edge case is explicitly covered, and no existing test assertions were weakened. The annotations being dropped for mixed text+citations responses is a deliberate, documented design choice. The only structural concern is Bedrock-specific field names in a general utility module, which is a maintainability note rather than a correctness issue.

litellm/litellm_core_utils/prompt_templates/common_utils.py — contains Bedrock-specific field names outside the llms/ tree; litellm/llms/bedrock/chat/converse_transformation.py — the annotation-stitching heuristic may need revisiting as Nova grounding response shapes evolve.

Important Files Changed

Filename	Overview
litellm/llms/bedrock/chat/converse_transformation.py	Adds _transform_citations_to_annotations to map Bedrock citationsContent blocks to OpenAI-style annotations, with conservative stitching that only attaches annotations when content is empty or punctuation-only. Logic and offset tracking look correct for the tested Nova grounding pattern.
litellm/litellm_core_utils/prompt_templates/factory.py	Refactors _convert_to_bedrock_tool_call_result into smaller helpers and adds _build_bedrock_tool_result_content_blocks for the search_results extension. Empty search_results correctly falls through to the normal content path.
litellm/litellm_core_utils/prompt_templates/common_utils.py	Adds extract_search_results_text with Bedrock-specific field names (source, title, content[].text, citations) in a general utility module, and calls it from convert_content_list_to_str for token estimation.
litellm/types/llms/bedrock.py	Adds SearchResultBlock, CitationSearchResultLocationBlock, CitationGeneratedContentBlock TypedDicts and extends existing types with new optional fields for search result and citation support.
litellm/types/llms/openai.py	Adds annotations field to ChatCompletionResponseMessage TypedDict; ChatCompletionAnnotation was already defined in the file.
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py	Filters custom_headers through the existing excluded_headers set before merging, preventing framework-injected defaults like content-length: 0 from overriding real passthrough response headers.
enterprise/litellm_enterprise/proxy/hooks/managed_files.py	Changes c["type"] to c.get("type") to avoid KeyError on Bedrock toolResult content blocks that omit the OpenAI-style type field.
litellm/litellm_core_utils/token_counter.py	Adds search_results branch in _count_messages to count tokens from the new search_results payload using extract_search_results_text.
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py	Adds comprehensive unit tests for search result mapping, citations-to-annotations conversion, offset tracking, and edge cases (empty search_results, null source/title, non-stitched annotations). All tests use mocked responses.

_{Reviews (6): Last reviewed commit: "fix(bedrock): complete citation TypedDic..." | Re-trigger Greptile}

Reformat the updated bedrock prompt template conversion file so CI black --check passes. Co-authored-by: Cursor <cursoragent@cursor.com>

cursor

Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.

^{Bugbot Autofix is ON, but it could not run because the spend limit has been reached. To enable Bugbot Autofix, have a team admin raise the spend limit in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 65f8e21. Configure here.}

veria-ai · 2026-05-28T08:05:50Z

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 2 · PR risk: 0/10

…nting Resolve mypy issues in citation parsing, only attach url_citation annotations when citation text is stitched into content, fall back to tool content when search_results is empty, and count search_results text in token/TPM preflight paths. Co-authored-by: Cursor <cursoragent@cursor.com>

Refactor _convert_to_bedrock_tool_call_result into smaller helpers so lint passes without changing Bedrock tool result behavior. Co-authored-by: Cursor <cursoragent@cursor.com>

…mates Include source, title, content text, and citations when estimating tokens so large metadata cannot bypass TPM preflight checks. Reformat factory.py with black. Co-authored-by: Cursor <cursoragent@cursor.com>

Sameerlite · 2026-05-29T13:57:17Z

@greptileai

…e_ids_from_messages

…itellm_fix_citation_bedrock

mateo-berri · 2026-05-30T02:43:43Z

@greptileai

mateo-berri · 2026-05-30T02:56:44Z

@greptileai

…blocks

mateo-berri · 2026-05-30T03:08:38Z

@greptileai

mateo-berri · 2026-05-30T03:26:28Z

@greptileai

mateo-berri

LGTM; thanks!

…cess tests/test_litellm is documented as mock-only (tests/test_litellm/readme.md), but the Bedrock passthrough regression test bound a real loopback socket and ran a live uvicorn server in a background thread. Drive the same finalizer output through an in-process ASGI server (httpx.ASGITransport) instead, asserting the recomputed Content-Length the client receives against the served body. The test still fails against the pre-#29120 leak, where the placeholder Response's content-length: 0 overrides the real upstream length, and passes once get_response_headers sanitizes custom headers, without opening a real socket.

…#29120) * Fix overiding of fastapi_response headers * fix(bedrock): support tool search results and surface citations as annotations Add an optional tool-message search_results path that maps directly to Bedrock toolResult.searchResult blocks, and convert Converse citationsContent into chat completion annotations for user-facing citation metadata. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(format): align bedrock prompt factory with black Reformat the updated bedrock prompt template conversion file so CI black --check passes. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(bedrock): harden citations, search_results mapping, and token counting Resolve mypy issues in citation parsing, only attach url_citation annotations when citation text is stitched into content, fall back to tool content when search_results is empty, and count search_results text in token/TPM preflight paths. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(bedrock): extract tool result helpers to satisfy PLR0915 Refactor _convert_to_bedrock_tool_call_result into smaller helpers so lint passes without changing Bedrock tool result behavior. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(bedrock): count all forwarded search_results fields in token estimates Include source, title, content text, and citations when estimating tokens so large metadata cannot bypass TPM preflight checks. Reformat factory.py with black. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(managed-files): skip content blocks without a type key in get_file_ids_from_messages * fix(bedrock): stitch citations for any punctuation-only text block * fix(bedrock): map null citation source/title to empty annotation strings * fix(bedrock): advance citation offset for text-only citationsContent blocks * fix(bedrock): complete citation TypedDicts for grounding annotations --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

Sameerlite and others added 2 commits May 28, 2026 12:20

Fix overiding of fastapi_response headers

adb10dc

greptile-apps Bot reviewed May 28, 2026

View reviewed changes

Comment thread litellm/litellm_core_utils/prompt_templates/factory.py Outdated

Comment thread litellm/llms/bedrock/chat/converse_transformation.py

fix(format): align bedrock prompt factory with black

65f8e21

Reformat the updated bedrock prompt template conversion file so CI black --check passes. Co-authored-by: Cursor <cursoragent@cursor.com>

cursor Bot reviewed May 28, 2026

View reviewed changes

Comment thread litellm/llms/bedrock/chat/converse_transformation.py

Comment thread litellm/litellm_core_utils/prompt_templates/factory.py Outdated

veria-ai Bot reviewed May 28, 2026

View reviewed changes

Comment thread litellm/litellm_core_utils/prompt_templates/factory.py Outdated

Sameerlite and others added 2 commits May 28, 2026 22:25

fix(bedrock): extract tool result helpers to satisfy PLR0915

dec4b3f

Refactor _convert_to_bedrock_tool_call_result into smaller helpers so lint passes without changing Bedrock tool result behavior. Co-authored-by: Cursor <cursoragent@cursor.com>

veria-ai Bot reviewed May 29, 2026

View reviewed changes

Comment thread litellm/litellm_core_utils/prompt_templates/factory.py

mateo-berri added 2 commits May 30, 2026 02:41

fix(managed-files): skip content blocks without a type key in get_fil…

c2bf793

…e_ids_from_messages

Merge remote-tracking branch 'origin/litellm_internal_staging' into l…

1e57e85

…itellm_fix_citation_bedrock

mateo-berri added 2 commits May 30, 2026 02:49

fix(bedrock): stitch citations for any punctuation-only text block

e279678

fix(bedrock): map null citation source/title to empty annotation strings

f870bc7

fix(bedrock): advance citation offset for text-only citationsContent …

2ea005b

…blocks

fix(bedrock): complete citation TypedDicts for grounding annotations

7aec1bd

mateo-berri approved these changes May 30, 2026

View reviewed changes

mateo-berri merged commit 1d9095f into litellm_internal_staging May 30, 2026
116 of 118 checks passed

mateo-berri deleted the litellm_fix_citation_bedrock branch May 30, 2026 03:48

mateo-berri mentioned this pull request Jun 3, 2026

test(passthrough): regression coverage for Bedrock Content-Length override #29549

Open

7 tasks

Uh oh!

Uh oh!

Conversation

Sameerlite commented May 28, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

codecov Bot commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

greptile-apps Bot commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

veria-ai Bot commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR overview

Security review

Uh oh!

Uh oh!

Sameerlite commented May 29, 2026

Uh oh!

mateo-berri commented May 30, 2026

Uh oh!

mateo-berri commented May 30, 2026

Uh oh!

mateo-berri commented May 30, 2026

Uh oh!

mateo-berri commented May 30, 2026

Uh oh!

mateo-berri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Sameerlite commented May 28, 2026 •

edited by cursor Bot

Loading

codecov Bot commented May 28, 2026 •

edited

Loading

greptile-apps Bot commented May 28, 2026 •

edited

Loading

veria-ai Bot commented May 28, 2026 •

edited

Loading