fix(bedrock): handle document content blocks with citations and context in Converse API by elvis-cai · Pull Request #26948 · BerriAI/litellm

elvis-cai · 2026-05-01T00:33:42Z

Relevant issues

Pre-Submission checklist

I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
My PR passes all unit tests on make test-unit
My PR's scope is as isolated as possible, it only solves 1 specific problem

Type

🐛 Bug Fix / ✨ Enhancement

Changes

document content blocks (PDF, CSV, etc. via Anthropic's API format) were silently dropped during Bedrock Converse API message conversion. Additionally, the Bedrock Converse API supports citations and context fields on DocumentBlock which were never forwarded.

Root cause

The for element in message_block["content"] loop in both the sync and async Bedrock Converse message paths had no elif element["type"] == "document" branch — blocks were skipped silently.
_convert_to_bedrock_tool_call_result() only handled text and image_url, dropping document blocks in tool results.
DocumentBlock in litellm/types/llms/bedrock.py was missing citations and context fields.

Fix

Added _process_document_message() static method on BedrockConverseMessagesProcessor:
- Validates source.type == "base64"
- Maps media_type → Bedrock format string (pdf, csv, docx, xlsx, html, txt, md, etc.)
- Generates a deterministic document name using SHA-256 content hashing
- Forwards citations: {enabled: true} → CitationsConfig(enabled=True) (Bedrock Converse API docs)
- Forwards context field
Added elif element["type"] == "document" branch in the async user-message path
Added elif element["type"] == "document" branch in the sync user-message path (_bedrock_converse_messages_pt)
Added elif content["type"] == "document" branch in _convert_to_bedrock_tool_call_result() (tool results)
Added CitationsConfig TypedDict and extended DocumentBlock with optional citations and context fields

Before / After

messages = [{"role": "user", "content": [
    {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "dGVzdA=="},
     "citations": {"enabled": True}, "context": "Q4 report"},
    {"type": "text", "text": "What is the title?"},
]}]

result = _bedrock_converse_messages_pt(messages, "anthropic.claude-sonnet-4-6", "bedrock")

# Before: [{'text': 'What is the title?'}]   — document dropped, citations/context lost
# After:  [{'document': {'source': {'bytes': 'dGVzdA=='}, 'format': 'pdf', 'name': '...',
#            'citations': {'enabled': True}, 'context': 'Q4 report'}},
#           {'text': 'What is the title?'}]

Tests added (`tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_document_blocks.py`)

Test	What it verifies
`test_document_block_basic`	Document block produces correct format, source.bytes, and name
`test_document_block_not_dropped_with_text`	Mixed document + text preserves both blocks
`test_document_block_citations_forwarded`	`citations.enabled=True` forwarded to `CitationsConfig`
`test_document_block_citations_not_set_when_disabled`	`citations` absent when disabled
`test_document_block_context_forwarded`	`context` field forwarded to DocumentBlock
`test_document_block_deterministic_name`	Same data always produces the same name
`test_document_format_mapping`	7 MIME types map to correct Bedrock format strings
`test_document_in_tool_result`	Document in tool result produces correct content block
`test_non_base64_source_raises`	Non-base64 source raises a clear error

🤖 Generated with Claude Code

CLAassistant · 2026-05-01T00:33:49Z

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

greptile-apps · 2026-05-01T00:36:16Z

Greptile Summary

This PR fixes silently-dropped document content blocks in the Bedrock Converse API message conversion path by adding a _process_document_message static method and elif branches in the sync, async, and tool-result loops. It also extends DocumentBlock with optional citations and context fields.

The test file imports _process_bedrock_document_block which does not exist in factory.py (the actual symbol is BedrockConverseMessagesProcessor._process_document_message), causing an ImportError that prevents all 9 tests from running.
test_non_base64_source_raises asserts ValueError, but the implementation raises litellm.BadRequestError, so that test will also fail independently.
DocumentBlock was changed to total=False, silently making the three Bedrock-required fields (format, source, name) optional at the type level.

Confidence Score: 2/5

The production fix logic is sound, but the test suite cannot run at all due to a broken import, leaving the fix unverified.

A P0 ImportError breaks every test added in this PR, and a separate P1 wrong-exception assertion would also fail independently. The DocumentBlock type regression is an additional P1. Multiple blocking issues across tests and types pull the score to 2.

test_bedrock_document_blocks.py (ImportError + wrong exception type) and litellm/types/llms/bedrock.py (DocumentBlock required-field regression)

Important Files Changed

Filename	Overview
litellm/litellm_core_utils/prompt_templates/factory.py	Adds `_process_document_message` static method and three `elif "document"` branches (async user path, sync user path, tool results path); MIME map is recreated per call (P2 style issue).
litellm/types/llms/bedrock.py	Adds `CitationsConfig` TypedDict and extends `DocumentBlock` — but changing `DocumentBlock` to `total=False` makes the three required Bedrock API fields (`format`, `source`, `name`) silently optional, undermining type safety.
tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_document_blocks.py	New test suite for document block handling — imports non-existent `_process_bedrock_document_block` (causes ImportError, breaking all 9 tests) and asserts `ValueError` where `BadRequestError` is actually raised.
tests/test_litellm/litellm_core_utils/prompt_templates/init.py	Empty `__init__.py` added to make the directory a proper Python package for test discovery — no issues.

_{Reviews (1): Last reviewed commit: "fix(bedrock): handle document content bl..." | Re-trigger Greptile}

greptile-apps · 2026-05-01T00:36:19Z

+from litellm.litellm_core_utils.prompt_templates.factory import (
+    _bedrock_converse_messages_pt,
+    _convert_to_bedrock_tool_call_result,
+    _process_bedrock_document_block,


Import error — _process_bedrock_document_block does not exist

The test file imports _process_bedrock_document_block from factory.py, but that name is never defined or exported there. The implementation lives as BedrockConverseMessagesProcessor._process_document_message. This causes an ImportError that prevents every test in this file from running.

greptile-apps · 2026-05-01T00:36:20Z

+def test_non_base64_source_raises():
+    """Non-base64 source type raises a clear ValueError."""
+    with pytest.raises(ValueError, match="base64"):
+        _process_bedrock_document_block(
+            {
+                "type": "document",
+                "source": {"type": "url", "url": "https://example.com/doc.pdf"},


Wrong exception type asserted

The implementation raises litellm.BadRequestError (which inherits from openai.BadRequestError → openai.APIStatusError → Exception, not ValueError). pytest.raises(ValueError) will not catch it, causing the test to fail with an unexpected BadRequestError.

greptile-apps · 2026-05-01T00:36:21Z

+class DocumentBlock(TypedDict, total=False):
    format: Union[BedrockDocumentTypes, str]
    source: SourceBlock
    name: str
+    citations: CitationsConfig
+    context: str


Required fields silently dropped from DocumentBlock

Changing DocumentBlock to total=False makes format, source, and name — which are all required by the Bedrock API — optional at the type level. Any caller that constructs a DocumentBlock without those fields will now pass type-checking but produce a runtime API error. Since only citations and context are genuinely optional, the required fields should use Required[...] annotations:

Suggested change

class DocumentBlock(TypedDict, total=False):

format: Union[BedrockDocumentTypes, str]

source: SourceBlock

name: str

citations: CitationsConfig

context: str

class DocumentBlock(TypedDict, total=False):

format: Required[Union[BedrockDocumentTypes, str]]

source: Required[SourceBlock]

name: Required[str]

citations: CitationsConfig

context: str

greptile-apps · 2026-05-01T00:36:23Z

+        _mime_to_format: dict = {
+            "application/pdf": "pdf",
+            "text/csv": "csv",
+            "application/msword": "doc",
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
+            "application/vnd.ms-excel": "xls",
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
+            "text/html": "html",
+            "text/plain": "txt",
+            "text/markdown": "md",
+        }
+        doc_format = _mime_to_format.get(media_type, media_type.split("/")[-1])


MIME map re-created on every call

_mime_to_format is a dict literal built inside a @staticmethod that can be called many times per request (once per document block). Promoting it to a module-level or class-level constant avoids repeated allocation with no correctness trade-off.

greptile-apps · 2026-05-01T00:36:23Z

+                if "document" in _doc_block:
+                    tool_result_content_blocks.append(
+                        BedrockToolResultContentBlock(document=_doc_block["document"])
+                    )

    message.get("name", "")


Dead expression — result of message.get("name", "") is discarded

This line is a no-op: the return value is never used or assigned. It should either be assigned (e.g. _tool_name = message.get("name", "")) or removed to avoid confusing future readers.

…xt in Converse API - Add `_process_document_message()` static method on `BedrockConverseMessagesProcessor` that converts Anthropic-style document blocks to Bedrock DocumentBlock format: - Validates source.type == "base64" - Maps media_type → Bedrock format string - Generates deterministic SHA-256-based document name - Forwards `citations.enabled` → CitationsConfig on the DocumentBlock - Forwards `context` field to DocumentBlock - Add `elif element["type"] == "document"` branch in async user-message path - Add `elif element["type"] == "document"` branch in sync user-message path - Add `elif content["type"] == "document"` branch in _convert_to_bedrock_tool_call_result() - Add CitationsConfig TypedDict and extend DocumentBlock with optional citations/context fields Fixes BerriAI#26937

codecov · 2026-05-01T00:46:17Z

Codecov Report

❌ Patch coverage is 93.18182% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...llm/litellm_core_utils/prompt_templates/factory.py	91.66%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

…and DocumentBlock required fields

github-actions · 2026-07-31T00:19:03Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs.

greptile-apps Bot reviewed May 1, 2026

View reviewed changes

elvis-cai force-pushed the feat/bedrock-document-citations-context branch from 4d65e7a to d038aae Compare May 1, 2026 00:43

fix: address review feedback — correct test imports, exception type, …

8411117

…and DocumentBlock required fields

stuxf mentioned this pull request May 1, 2026

chore(proxy): tighten router-settings-override and mock-testing trust #26968

Merged

4 tasks

github-actions Bot added the stale label Jul 31, 2026

github-actions Bot closed this Aug 7, 2026

Navigation Menu

Uh oh!

Uh oh!

fix(bedrock): handle document content blocks with citations and context in Converse API - #26948

fix(bedrock): handle document content blocks with citations and context in Converse API#26948
elvis-cai wants to merge 2 commits into
BerriAI:litellm_internal_stagingfrom
elvis-cai:feat/bedrock-document-citations-context

elvis-cai commented May 1, 2026

Uh oh!

CLAassistant commented May 1, 2026

Uh oh!

greptile-apps Bot commented May 1, 2026

Important Files Changed

Uh oh!

greptile-apps Bot May 1, 2026

Uh oh!

greptile-apps Bot May 1, 2026

Uh oh!

greptile-apps Bot May 1, 2026

Uh oh!

greptile-apps Bot May 1, 2026

Uh oh!

greptile-apps Bot May 1, 2026

Uh oh!

codecov Bot commented May 1, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

elvis-cai commented May 1, 2026

Relevant issues

Pre-Submission checklist

Type

Changes

Root cause

Fix

Before / After

Tests added (tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_document_blocks.py)

Uh oh!

CLAassistant commented May 1, 2026

Uh oh!

greptile-apps Bot commented May 1, 2026

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Uh oh!

greptile-apps Bot May 1, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot May 1, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot May 1, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot May 1, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot May 1, 2026

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Tests added (`tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_document_blocks.py`)

codecov Bot commented May 1, 2026 •

edited

Loading