diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d29ca1649ff..31612b8e78d 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3984,6 +3984,14 @@ def _convert_to_bedrock_tool_call_result( tool_result_content_blocks.append( BedrockToolResultContentBlock(image=_block["image"]) ) + elif content["type"] == "document": + _doc_block = BedrockConverseMessagesProcessor._process_document_message( + content + ) + if "document" in _doc_block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(document=_doc_block["document"]) + ) message.get("name", "") id = str(message.get("tool_call_id", str(uuid.uuid4()))) @@ -4439,6 +4447,11 @@ class BedrockConverseMessagesProcessor: message=cast(ChatCompletionFileObject, element) ) _parts.append(_part) + elif element["type"] == "document": + _part = BedrockConverseMessagesProcessor._process_document_message( + element + ) + _parts.append(_part) _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( message_block=cast( @@ -4718,6 +4731,36 @@ class BedrockConverseMessagesProcessor: image_url=cast(str, file_id or file_data), format=format ) + @staticmethod + def _process_document_message(element: dict) -> BedrockContentBlock: + """Convert a document content block to a Bedrock DocumentBlock. + + Handles the Anthropic-style document format: + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "..."}} + """ + source = element["source"] + media_type: str = source["media_type"] + data: str = source["data"] + doc_format = media_type.split("/")[1] + + # Deterministic name using the same hashing pattern as _create_bedrock_block + HASH_SAMPLE_BYTES = 64 * 1024 + normalized = "".join(data.split()).encode("utf-8") + sample = normalized[:HASH_SAMPLE_BYTES] + hasher = hashlib.sha256() + hasher.update(sample) + hasher.update(str(len(normalized)).encode("utf-8")) + content_hash = hasher.hexdigest()[:16] + document_name = f"Document_{content_hash}_{doc_format}" + + return BedrockContentBlock( + document=BedrockDocumentBlock( + source=BedrockSourceBlock(bytes=data), + format=doc_format, + name=document_name, + ) + ) + @staticmethod def add_thinking_blocks_to_assistant_content( thinking_blocks: List[BedrockContentBlock], @@ -4815,6 +4858,11 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 ) ) _parts.append(_part) + elif element["type"] == "document": + _part = BedrockConverseMessagesProcessor._process_document_message( + element + ) + _parts.append(_part) _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( message_block=cast( diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 30b47a853ef..d6005f25da9 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -9,8 +9,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, BedrockConverseMessagesProcessor, BedrockImageProcessor, - anthropic_messages_pt, + _bedrock_converse_messages_pt, _convert_to_bedrock_tool_call_invoke, + _convert_to_bedrock_tool_call_result, + anthropic_messages_pt, convert_to_gemini_tool_call_result, ollama_pt, sanitize_messages_for_tool_calling, @@ -2339,3 +2341,168 @@ def test_anthropic_messages_pt_file_block_preserves_cache_control(): assert text_block["type"] == "text" assert "cache_control" in text_block assert text_block["cache_control"]["type"] == "ephemeral" + + +def test_bedrock_converse_messages_pt_document_block(): + """Test that a document content block is converted to a Bedrock DocumentBlock.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + }, + ], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + assert len(result) == 1 + content = result[0]["content"] + assert len(content) == 1 + + doc_block = content[0] + assert "document" in doc_block + assert doc_block["document"]["format"] == "pdf" + assert doc_block["document"]["source"]["bytes"] == "dGVzdA==" + assert doc_block["document"]["name"].startswith("Document_") + assert doc_block["document"]["name"].endswith("_pdf") + + +def test_bedrock_converse_messages_pt_document_and_text(): + """Test that mixed document + text content produces both blocks.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + }, + {"type": "text", "text": "What is the title?"}, + ], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + content = result[0]["content"] + assert len(content) == 2 + + doc_block = content[0] + assert "document" in doc_block + assert doc_block["document"]["format"] == "pdf" + + text_block = content[1] + assert "text" in text_block + assert text_block["text"] == "What is the title?" + + +def test_bedrock_converse_messages_pt_document_in_tool_result(): + """Test that a tool result containing a document block is converted correctly.""" + message = { + "role": "tool", + "tool_call_id": "tool_123", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + }, + ], + } + + result = _convert_to_bedrock_tool_call_result(message) + + tool_result = result["toolResult"] + assert len(tool_result["content"]) == 1 + assert "document" in tool_result["content"][0] + assert tool_result["content"][0]["document"]["format"] == "pdf" + assert tool_result["content"][0]["document"]["source"]["bytes"] == "dGVzdA==" + + +def test_bedrock_converse_messages_pt_document_various_formats(): + """Test that various document media types produce the correct format value.""" + test_cases = [ + ("application/pdf", "pdf"), + ("text/csv", "csv"), + ("text/html", "html"), + ( + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "vnd.openxmlformats-officedocument.wordprocessingml.document", + ), + ] + + for media_type, expected_format in test_cases: + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": media_type, + "data": "dGVzdA==", + }, + }, + ], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + doc_block = result[0]["content"][0] + assert doc_block["document"]["format"] == expected_format, ( + f"Expected format '{expected_format}' for media_type '{media_type}', " + f"got '{doc_block['document']['format']}'" + ) + + +def test_bedrock_converse_messages_pt_document_deterministic_name(): + """Test that the same document data always produces the same name.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + }, + ], + } + ] + + result1 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + result2 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + name1 = result1[0]["content"][0]["document"]["name"] + name2 = result2[0]["content"][0]["document"]["name"] + assert name1 == name2