From 5bba1e84050fd603bcb2dbff8cb590fbc6bf4777 Mon Sep 17 00:00:00 2001 From: langping Date: Wed, 29 Oct 2025 23:24:20 +0800 Subject: [PATCH] Added fallback logic for detecting file content-type when S3 returns generic (#15635) * enhance image processing fallback logic * Extract to comment utils --- .../prompt_templates/common_utils.py | 96 +++++++ .../prompt_templates/factory.py | 25 +- ...llm_core_utils_prompt_templates_factory.py | 253 +++++++++++++++++- 3 files changed, 365 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 19d5932ff2..33658d4906 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -654,6 +654,102 @@ def _get_image_mime_type_from_url(url: str) -> Optional[str]: return None +def infer_content_type_from_url_and_content( + url: str, + content: bytes, + current_content_type: Optional[str] = None, +) -> str: + """ + Infer content type from URL extension and binary content when content-type header is missing or generic. + + This helper implements a fallback strategy for determining MIME types when HTTP headers + are missing or provide generic values (like binary/octet-stream). It's commonly used + when processing images and documents from various sources (S3, URLs, etc.). + + Fallback Strategy: + 1. If current_content_type is valid (not None and not generic octet-stream), return it + 2. Try to infer from URL extension (handles query parameters) + 3. Try to detect from binary content signature (magic bytes) + 4. Raise ValueError if all methods fail + + Args: + url: The URL of the content (used to extract file extension) + content: The binary content (first ~100 bytes are sufficient for detection) + current_content_type: The current content-type from headers (may be None or generic) + + Returns: + str: The inferred MIME type (e.g., "image/png", "application/pdf") + + Raises: + ValueError: If content type cannot be determined by any method + + Example: + >>> content_type = infer_content_type_from_url_and_content( + ... url="https://s3.amazonaws.com/bucket/image.png?AWSAccessKeyId=123", + ... content=png_binary_data, + ... current_content_type="binary/octet-stream" + ... ) + >>> print(content_type) + "image/png" + """ + from litellm.litellm_core_utils.token_counter import get_image_type + + # If we have a valid content type that's not generic, use it + if current_content_type and current_content_type not in [ + "binary/octet-stream", + "application/octet-stream", + ]: + return current_content_type + + # Extension to MIME type mapping + # Supports images, documents, and other common file types + extension_to_mime = { + # Image formats + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "png": "image/png", + "gif": "image/gif", + "webp": "image/webp", + # Document formats + "pdf": "application/pdf", + "csv": "text/csv", + "doc": "application/msword", + "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "xls": "application/vnd.ms-excel", + "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "html": "text/html", + "txt": "text/plain", + "md": "text/markdown", + } + + # Try to infer from URL extension + if url: + extension = url.split(".")[-1].lower().split("?")[0] # Remove query params + inferred_type = extension_to_mime.get(extension) + if inferred_type: + return inferred_type + + # Try to detect from binary content signature (magic bytes) + if content: + detected_type = get_image_type(content[:100]) + if detected_type: + type_to_mime = { + "png": "image/png", + "jpeg": "image/jpeg", + "gif": "image/gif", + "webp": "image/webp", + "heic": "image/heic", + } + if detected_type in type_to_mime: + return type_to_mime[detected_type] + + # If all fallbacks failed, raise error + raise ValueError( + f"Unable to determine content type from URL: {url}. " + f"Response content-type: {current_content_type}" + ) + + def get_tool_call_names(tools: List[ChatCompletionToolParam]) -> List[str]: """ Get tool call names from tools diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index a5a4ab6cbf..337298688e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -14,6 +14,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client +from litellm.litellm_core_utils.token_counter import get_image_type from litellm.types.files import get_file_extension_from_mime_type from litellm.types.llms.anthropic import * from litellm.types.llms.bedrock import CachePointBlock @@ -38,7 +39,11 @@ from litellm.types.llms.vertex_ai import FunctionResponse as VertexFunctionRespo from litellm.types.llms.vertex_ai import PartType as VertexPartType from litellm.types.utils import GenericImageParsingChunk -from .common_utils import convert_content_list_to_str, is_non_content_values_set +from .common_utils import ( + convert_content_list_to_str, + infer_content_type_from_url_and_content, + is_non_content_values_set, +) from .image_handling import convert_url_to_base64 @@ -2536,13 +2541,17 @@ class BedrockImageProcessor: """Handles both sync and async image processing for Bedrock conversations.""" @staticmethod - def _post_call_image_processing(response: httpx.Response) -> Tuple[str, str]: + def _post_call_image_processing(response: httpx.Response, image_url: str = "") -> Tuple[str, str]: # Check the response's content type to ensure it is an image content_type = response.headers.get("content-type") - if not content_type: - raise ValueError( - f"URL does not contain content-type (content-type: {content_type})" - ) + + # Use helper function to infer content type with fallback logic + content_type = infer_content_type_from_url_and_content( + url=image_url, + content=response.content, + current_content_type=content_type, + ) + content_type = _parse_content_type(content_type) # Convert the image content to base64 bytes @@ -2561,7 +2570,7 @@ class BedrockImageProcessor: response = await client.get(image_url, follow_redirects=True) response.raise_for_status() # Raise an exception for HTTP errors - return BedrockImageProcessor._post_call_image_processing(response) + return BedrockImageProcessor._post_call_image_processing(response, image_url) except Exception as e: raise e @@ -2574,7 +2583,7 @@ class BedrockImageProcessor: response = client.get(image_url, follow_redirects=True) response.raise_for_status() # Raise an exception for HTTP errors - return BedrockImageProcessor._post_call_image_processing(response) + return BedrockImageProcessor._post_call_image_processing(response, image_url) except Exception as e: raise e 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 a40c94af1d..c23b4f98df 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 @@ -1,5 +1,5 @@ import json -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -567,6 +567,257 @@ def test_bedrock_tools_unpack_defs(): _bedrock_tools_pt(tools=tools) +def test_bedrock_image_processor_content_type_fallback_url_extension(): + """ + Test that _post_call_image_processing falls back to URL extension + when content-type is binary/octet-stream or application/octet-stream + """ + import base64 + + # Create mock response with binary/octet-stream content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = "binary/octet-stream" + + # Create a simple PNG header (magic bytes) + png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" + png_content = png_header + b"\x00" * 100 # Add some padding + mock_response.content = png_content + + # Test with .png URL + image_url = "https://example.com/test-image.png" + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url + ) + + assert content_type == "image/png" + assert base64_bytes == base64.b64encode(png_content).decode("utf-8") + + +def test_bedrock_image_processor_content_type_fallback_binary_detection(): + """ + Test that _post_call_image_processing falls back to binary content detection + when content-type is missing and URL extension is not recognized + """ + import base64 + + # Create mock response with no content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = None + + # Create a JPEG header (magic bytes) + jpeg_header = b"\xff\xd8\xff" + jpeg_content = jpeg_header + b"\x00" * 100 # Add some padding + mock_response.content = jpeg_content + + # Test with URL without extension + image_url = "https://example.com/test-image-without-extension" + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url + ) + + assert content_type == "image/jpeg" + assert base64_bytes == base64.b64encode(jpeg_content).decode("utf-8") + + +def test_bedrock_image_processor_content_type_fallback_application_octet_stream(): + """ + Test that _post_call_image_processing handles application/octet-stream correctly + """ + import base64 + + # Create mock response with application/octet-stream content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = "application/octet-stream" + + # Create a GIF header (magic bytes) + gif_header = b"GIF8" + b"\x00" + b"a" + gif_content = gif_header + b"\x00" * 100 # Add some padding + mock_response.content = gif_content + + # Test with .gif URL + image_url = "https://s3.amazonaws.com/bucket/image.gif" + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url + ) + + assert content_type == "image/gif" + assert base64_bytes == base64.b64encode(gif_content).decode("utf-8") + + +def test_bedrock_image_processor_content_type_with_query_params(): + """ + Test that _post_call_image_processing correctly extracts extension from URL with query parameters + """ + import base64 + + # Create mock response with binary/octet-stream content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = "binary/octet-stream" + + # Create a WebP header (magic bytes) + webp_header = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" + webp_content = webp_header + b"\x00" * 100 # Add some padding + mock_response.content = webp_content + + # Test with URL containing query parameters (common in S3 signed URLs) + image_url = "https://s3.amazonaws.com/bucket/image.webp?AWSAccessKeyId=123&Expires=456&Signature=789" + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url + ) + + assert content_type == "image/webp" + assert base64_bytes == base64.b64encode(webp_content).decode("utf-8") + + +def test_bedrock_image_processor_content_type_normal_header(): + """ + Test that _post_call_image_processing works normally when content-type is correctly set + """ + import base64 + + # Create mock response with correct content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = "image/png" + + # Create a PNG header + png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" + png_content = png_header + b"\x00" * 100 + mock_response.content = png_content + + image_url = "https://example.com/test-image.png" + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url + ) + + assert content_type == "image/png" + assert base64_bytes == base64.b64encode(png_content).decode("utf-8") + + +def test_bedrock_image_processor_content_type_fallback_failure(): + """ + Test that _post_call_image_processing raises ValueError when all fallback methods fail + """ + # Create mock response with binary/octet-stream content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = "binary/octet-stream" + + # Create content with unrecognizable image format + mock_response.content = b"\x00" * 100 + + # Test with URL without recognizable extension + image_url = "https://example.com/unknown-file" + + with pytest.raises(ValueError) as excinfo: + BedrockImageProcessor._post_call_image_processing(mock_response, image_url) + + assert "Unable to determine content type" in str(excinfo.value) + + +def test_bedrock_image_processor_content_type_jpeg_variants(): + """ + Test that _post_call_image_processing handles both .jpg and .jpeg extensions correctly + """ + # Create mock response with binary/octet-stream + mock_response = MagicMock() + mock_response.headers.get.return_value = "binary/octet-stream" + + jpeg_header = b"\xff\xd8\xff" + jpeg_content = jpeg_header + b"\x00" * 100 + mock_response.content = jpeg_content + + # Test with .jpg extension + image_url_jpg = "https://example.com/photo.jpg" + _, content_type_jpg = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url_jpg + ) + assert content_type_jpg == "image/jpeg" + + # Test with .jpeg extension + image_url_jpeg = "https://example.com/photo.jpeg" + _, content_type_jpeg = BedrockImageProcessor._post_call_image_processing( + mock_response, image_url_jpeg + ) + assert content_type_jpeg == "image/jpeg" + + +def test_bedrock_image_processor_content_type_pdf_document(): + """ + Test that _post_call_image_processing handles PDF documents correctly + when content-type is binary/octet-stream + """ + import base64 + + # Create mock response with binary/octet-stream content-type + mock_response = MagicMock() + mock_response.headers.get.return_value = "binary/octet-stream" + + # Create a PDF header (magic bytes: %PDF) + pdf_header = b"%PDF-1.4" + pdf_content = pdf_header + b"\x00" * 100 + mock_response.content = pdf_content + + # Test with .pdf URL + pdf_url = "https://s3.amazonaws.com/bucket/document.pdf" + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, pdf_url + ) + + assert content_type == "application/pdf" + assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") + + +def test_bedrock_image_processor_content_type_document_formats(): + """ + Test that _post_call_image_processing handles various document formats + """ + import base64 + + # Create mock response + mock_response = MagicMock() + mock_response.headers.get.return_value = "application/octet-stream" + mock_response.content = b"\x00" * 100 + + # Test various document formats + test_cases = [ + ("https://example.com/doc.pdf", "application/pdf"), + ("https://example.com/sheet.csv", "text/csv"), + ("https://example.com/doc.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"), + ("https://example.com/sheet.xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"), + ("https://example.com/page.html", "text/html"), + ("https://example.com/readme.txt", "text/plain"), + ] + + for url, expected_mime in test_cases: + _, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, url + ) + assert content_type == expected_mime, f"Expected {expected_mime} for {url}, got {content_type}" + + +def test_bedrock_image_processor_content_type_s3_pdf_with_query(): + """ + Test that _post_call_image_processing handles S3 PDF with query parameters + """ + import base64 + + # Create mock response + mock_response = MagicMock() + mock_response.headers.get.return_value = "binary/octet-stream" + + pdf_content = b"%PDF-1.4" + b"\x00" * 100 + mock_response.content = pdf_content + + # S3 signed URL with query parameters + s3_url = "https://my-bucket.s3.us-east-1.amazonaws.com/documents/report.pdf?AWSAccessKeyId=AKIAIOSFODNN7EXAMPLE&Expires=1234567890&Signature=abcdef123456" + + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( + mock_response, s3_url + ) + + assert content_type == "application/pdf" + assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") + + def test_bedrock_tools_pt_empty_description(): """ Test that _bedrock_tools_pt handles empty string descriptions correctly.