From 2c9970742ff8a593800afa67f9fe3eb407681e99 Mon Sep 17 00:00:00 2001 From: Jamie Goodyear Date: Wed, 5 Nov 2025 23:09:38 -0330 Subject: [PATCH] [LiteLLM-16250] Proxy to Bedrock will add name to file content, breaks when cache_control in use (#16275) --- .../prompt_templates/factory.py | 30 ++- ...llm_core_utils_prompt_templates_factory.py | 229 ++++++++++++++++++ 2 files changed, 258 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 1dd4590a0f..717c260765 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1,4 +1,5 @@ import copy +import hashlib import json import mimetypes import re @@ -2706,12 +2707,39 @@ class BedrockImageProcessor: for video_type in supported_video_formats ) + HASH_SAMPLE_BYTES = 64 * 1024 # hash up to 64 KB of data + if is_document: + # --- Prepare normalized bytes for hashing (without modifying original) --- + if isinstance(image_bytes, str): + # Remove whitespace/newlines so base64 variations hash identically + normalized = "".join(image_bytes.split()).encode("utf-8") + else: + normalized = image_bytes + + # --- Use only the first 64 KB for speed --- + if len(normalized) <= HASH_SAMPLE_BYTES: + sample = normalized + else: + sample = normalized[:HASH_SAMPLE_BYTES] + + # --- Compute deterministic hash (sample + total length) --- + hasher = hashlib.sha256() + hasher.update(sample) + hasher.update( + str(len(normalized)).encode("utf-8") + ) # include full length for uniqueness + full_hash = hasher.hexdigest() + content_hash = full_hash[:16] # short deterministic ID + + document_name = f"DocumentPDFmessages_{content_hash}_{image_format}" + + # --- Return content block --- return BedrockContentBlock( document=BedrockDocumentBlock( source=_blob, format=image_format, - name=f"DocumentPDFmessages_{str(uuid.uuid4())}", + name=document_name, ) ) elif is_video: 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 c23b4f98df..41ac893b4d 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 @@ -857,3 +857,232 @@ def test_bedrock_tools_pt_empty_description(): assert tool_spec is not None assert tool_spec.get("name") == "get_weather" assert tool_spec.get("description") == "get_weather" + +def test_bedrock_create_bedrock_block_deterministic_document_hash(): + """ + Test that _create_bedrock_block generates deterministic document names + based on content hash. Same content should produce same hash. + """ + import base64 + + # Create PDF content + pdf_content = b"%PDF-1.4\nSome PDF content here" + b"\x00" * 100 + base64_content = base64.b64encode(pdf_content).decode("utf-8") + + # Create two blocks with same content + block1 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + ) + block2 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + ) + + # Both should have the same document name + assert block1.get("document") is not None + assert block2.get("document") is not None + assert block1["document"]["name"] == block2["document"]["name"] + assert "DocumentPDFmessages_" in block1["document"]["name"] + + +def test_bedrock_create_bedrock_block_different_content_different_hash(): + """ + Test that different content produces different document hashes. + """ + import base64 + + # Create two different PDF contents + pdf_content1 = b"%PDF-1.4\nFirst PDF content" + b"\x00" * 100 + pdf_content2 = b"%PDF-1.4\nSecond PDF content" + b"\x00" * 100 + + base64_content1 = base64.b64encode(pdf_content1).decode("utf-8") + base64_content2 = base64.b64encode(pdf_content2).decode("utf-8") + + # Create blocks + block1 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content1, mime_type="application/pdf", image_format="pdf" + ) + block2 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content2, mime_type="application/pdf", image_format="pdf" + ) + + # Should have different document names + assert block1["document"]["name"] != block2["document"]["name"] + + +def test_bedrock_create_bedrock_block_normalized_base64(): + """ + Test that different base64 formatting (with/without whitespace) + produces the same hash due to normalization. + """ + import base64 + + pdf_content = b"%PDF-1.4\nTest content" + b"\x00" * 100 + base64_content = base64.b64encode(pdf_content).decode("utf-8") + + # Create versions with different whitespace + base64_with_newlines = "\n".join( + [base64_content[i : i + 64] for i in range(0, len(base64_content), 64)] + ) + base64_with_spaces = " ".join( + [base64_content[i : i + 32] for i in range(0, len(base64_content), 32)] + ) + + # Create blocks + block1 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + ) + block2 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_with_newlines, + mime_type="application/pdf", + image_format="pdf", + ) + block3 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_with_spaces, + mime_type="application/pdf", + image_format="pdf", + ) + + # All should have the same document name due to normalization + assert block1["document"]["name"] == block2["document"]["name"] + assert block1["document"]["name"] == block3["document"]["name"] + + +def test_bedrock_create_bedrock_block_large_file_sampling(): + """ + Test that files larger than 64KB use sampling correctly and + different lengths produce different hashes. + """ + import base64 + + # Create two large files with same first 64KB but different total lengths + first_64kb = b"%PDF-1.4\n" + b"A" * (64 * 1024) + large_content1 = first_64kb + b"X" * 1024 # 65KB + large_content2 = first_64kb + b"Y" * 2048 # 66KB + + base64_content1 = base64.b64encode(large_content1).decode("utf-8") + base64_content2 = base64.b64encode(large_content2).decode("utf-8") + + # Create blocks + block1 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content1, mime_type="application/pdf", image_format="pdf" + ) + block2 = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content2, mime_type="application/pdf", image_format="pdf" + ) + + # Should have different names because total length is different + assert block1["document"]["name"] != block2["document"]["name"] + + +def test_bedrock_create_bedrock_block_very_large_file(): + """ + Test that very large files (>64KB) are handled correctly. + """ + import base64 + + # Create a large file (100KB) + large_content = b"%PDF-1.4\n" + b"X" * (100 * 1024) + base64_content = base64.b64encode(large_content).decode("utf-8") + + # Create block + block = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + ) + + # Should have a valid document name + assert block.get("document") is not None + assert "DocumentPDFmessages_" in block["document"]["name"] + assert block["document"]["format"] == "pdf" + + +def test_bedrock_create_bedrock_block_image_type(): + """ + Test that image types still work correctly (no document name). + """ + import base64 + + # Create PNG content + png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" + png_content = png_header + b"\x00" * 100 + base64_content = base64.b64encode(png_content).decode("utf-8") + + # Create block + block = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="image/png", image_format="png" + ) + + # Should be an image block, not a document block + assert block.get("image") is not None + assert block.get("document") is None + assert block["image"]["format"] == "png" + + +def test_bedrock_create_bedrock_block_video_type(): + """ + Test that video types still work correctly (no document name). + """ + import base64 + + # Create MP4 content + mp4_content = b"\x00\x00\x00\x20\x66\x74\x79\x70" + b"\x00" * 100 + base64_content = base64.b64encode(mp4_content).decode("utf-8") + + # Create block + block = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="video/mp4", image_format="mp4" + ) + + # Should be a video block, not a document block + assert block.get("video") is not None + assert block.get("document") is None + assert block["video"]["format"] == "mp4" + + +def test_bedrock_create_bedrock_block_document_name_format(): + """ + Test that document names follow the expected format: + DocumentPDFmessages_{16_char_hash}_{format} + """ + import base64 + import re + + pdf_content = b"%PDF-1.4\nTest content" + b"\x00" * 100 + base64_content = base64.b64encode(pdf_content).decode("utf-8") + + block = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + ) + + document_name = block["document"]["name"] + + # Check format: DocumentPDFmessages_{16_hex_chars}_{format} + pattern = r"^DocumentPDFmessages_[0-9a-f]{16}_pdf$" + assert re.match(pattern, document_name), f"Document name format mismatch: {document_name}" + + +def test_bedrock_create_bedrock_block_different_document_formats(): + """ + Test that different document formats (PDF, CSV, DOCX) are handled correctly. + """ + import base64 + + test_cases = [ + (b"%PDF-1.4\nContent", "application/pdf", "pdf"), + (b"col1,col2\nval1,val2", "text/csv", "csv"), + ( + b"PK\x03\x04", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "docx", + ), + ] + + for content, mime_type, format_type in test_cases: + base64_content = base64.b64encode(content + b"\x00" * 100).decode("utf-8") + block = BedrockImageProcessor._create_bedrock_block( + image_bytes=base64_content, mime_type=mime_type, image_format=format_type + ) + + assert block.get("document") is not None + assert f"DocumentPDFmessages_" in block["document"]["name"] + assert block["document"]["name"].endswith(f"_{format_type}") + assert block["document"]["format"] == format_type