From e12e242b4603c59f0985fc4551b86772e4a90d76 Mon Sep 17 00:00:00 2001 From: Jerry-Xin Date: Sat, 7 Mar 2026 17:28:39 +0800 Subject: [PATCH] fix(vertex_ai): handle invalid Content-Type for image URLs Addresses PR #23026 review feedback: 1. image_handling.py - _process_image_response: - Strip Content-Type parameters (e.g., 'image/png; charset=utf-8' -> 'image/png') - Validate Content-Type against supported types (image/jpeg, png, gif, webp) - Fall back to URL extension when Content-Type is missing OR invalid - Handle signed URLs by stripping query params before extracting extension 2. transformation.py - _convert_image_urls_to_base64: - Use convert_to_anthropic_image_obj() instead of manual data URL parsing - Use {\*\*block, 'source': ...} to preserve all original block fields This fixes the issue where CDNs like Aliyun OSS return incorrect Content-Type (e.g., application/x-www-form-urlencoded for .png files), causing Vertex AI to reject with invalid media_type error. --- .../prompt_templates/image_handling.py | 75 ++++-- .../transformation.py | 44 ++-- .../litellm_core_utils/test_image_handling.py | 229 +++++++++++++++++- ..._vertex_ai_anthropic_image_url_handling.py | 48 ++-- 4 files changed, 317 insertions(+), 79 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 7137a4e4222..c7169bf17e1 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -16,6 +16,56 @@ MAX_IMGS_IN_MEMORY = 10 in_memory_cache = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) +# Supported image media types for Anthropic/Vertex AI +SUPPORTED_IMAGE_TYPES = {"image/jpeg", "image/png", "image/gif", "image/webp"} + +# Mapping from file extensions to media types +EXTENSION_TO_MEDIA_TYPE = { + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "png": "image/png", + "gif": "image/gif", + "webp": "image/webp", +} + + +def _infer_media_type_from_url(url: str) -> str: + """ + Infer media type from URL extension. + Handles signed URLs by stripping query parameters before extracting extension. + """ + # Strip query parameters for signed URLs (e.g., ?x-oss-signature=...) + url_without_query = url.split("?")[0] + extension = url_without_query.split(".")[-1].lower() + media_type = EXTENSION_TO_MEDIA_TYPE.get(extension) + if media_type is None: + raise Exception( + f"Error: Unsupported image format. Extension={extension}. " + f"Supported types = {list(SUPPORTED_IMAGE_TYPES)}" + ) + return media_type + + +def _get_valid_media_type(content_type: str | None, url: str) -> str: + """ + Get a valid media type from Content-Type header, falling back to URL extension. + + - Strips Content-Type parameters (e.g., 'image/png; charset=utf-8' -> 'image/png') + - Validates against supported types + - Falls back to URL extension inference if Content-Type is missing or invalid + """ + if content_type is not None: + # Strip parameters (e.g., "image/png; charset=utf-8" -> "image/png") + media_type = content_type.split(";")[0].strip() + if media_type in SUPPORTED_IMAGE_TYPES: + return media_type + # Content-Type is invalid (e.g., application/octet-stream, application/x-www-form-urlencoded) + # Fall through to URL extension inference + + # Fallback to URL extension + return _infer_media_type_from_url(url) + + def _process_image_response(response: Response, url: str) -> str: if response.status_code != 200: raise litellm.ImageFetchError( @@ -35,7 +85,7 @@ def _process_image_response(response: Response, url: str) -> str: max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) image_bytes = bytearray() bytes_downloaded = 0 - + for chunk in response.iter_bytes(chunk_size=8192): bytes_downloaded += len(chunk) if bytes_downloaded > max_bytes: @@ -44,28 +94,13 @@ def _process_image_response(response: Response, url: str) -> str: f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) image_bytes.extend(chunk) - + base64_image = base64.b64encode(image_bytes).decode("utf-8") - image_type = response.headers.get("Content-Type") - if image_type is None: - img_type = url.split(".")[-1].lower() - _img_type = { - "jpg": "image/jpeg", - "jpeg": "image/jpeg", - "png": "image/png", - "gif": "image/gif", - "webp": "image/webp", - }.get(img_type) - if _img_type is None: - raise Exception( - f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" - ) - img_type = _img_type - else: - img_type = image_type + content_type = response.headers.get("Content-Type") + media_type = _get_valid_media_type(content_type, url) - result = f"data:{img_type};base64,{base64_image}" + result = f"data:{media_type};base64,{base64_image}" in_memory_cache.set_cache(url, result) return result diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index b15000e4db8..652905649d4 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,7 +1,7 @@ from typing import Any, Dict, List, Optional, Tuple -from litellm.litellm_core_utils.prompt_templates.image_handling import ( - convert_url_to_base64, +from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_image_obj, ) from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( @@ -171,30 +171,22 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert if isinstance(source, dict) and source.get("type") == "url": url = source.get("url") if url: - # Convert URL to base64 - base64_data_url = convert_url_to_base64(url=url) - # Parse the data URL: data:image/jpeg;base64, - if base64_data_url.startswith("data:"): - # Extract media type and data - parts = base64_data_url.split(";base64,", 1) - if len(parts) == 2: - media_type = parts[0].replace("data:", "") - data = parts[1] - new_block = { - "type": "image", - "source": { - "type": "base64", - "media_type": media_type, - "data": data, - }, - } - # Preserve cache_control if present - if "cache_control" in block: - new_block["cache_control"] = block[ - "cache_control" - ] - new_content.append(new_block) - continue + # Convert URL to base64 using existing utility + image_obj = convert_to_anthropic_image_obj( + openai_image_url=url, format=None + ) + # Preserve all original block fields (e.g., cache_control) + # while replacing the source + new_block = { + **block, + "source": { + "type": image_obj["type"], + "media_type": image_obj["media_type"], + "data": image_obj["data"], + }, + } + new_content.append(new_block) + continue new_content.append(block) diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 9c2939b2da5..cf5ab5b69c5 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -6,6 +6,8 @@ from httpx import Request, Response import litellm from litellm import constants from litellm.litellm_core_utils.prompt_templates.image_handling import ( + _get_valid_media_type, + _infer_media_type_from_url, convert_url_to_base64, ) @@ -37,9 +39,7 @@ def test_completion_with_invalid_image_url(monkeypatch): } ] with pytest.raises(litellm.ImageFetchError) as excinfo: - litellm.completion( - model="gemini/gemini-pro", messages=messages, api_key="test" - ) + litellm.completion(model="gemini/gemini-pro", messages=messages, api_key="test") assert excinfo.value.status_code == 400 assert "Unable to fetch image" in str(excinfo.value) @@ -81,7 +81,7 @@ class StreamingLargeImageClient: headers = {"Content-Type": "image/jpeg"} if self.include_content_length: headers["Content-Length"] = str(size_bytes) - + # Create a generator that yields chunks without creating the whole file in memory def generate_chunks(total_size, chunk_size=8192): bytes_sent = 0 @@ -89,7 +89,7 @@ class StreamingLargeImageClient: chunk = b"x" * min(chunk_size, total_size - bytes_sent) bytes_sent += len(chunk) yield chunk - + # Create response with streaming content response = Response( status_code=200, @@ -97,7 +97,9 @@ class StreamingLargeImageClient: request=Request("GET", url), ) # Mock the iter_bytes method to return our generator - response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) + response.iter_bytes = lambda chunk_size=8192: generate_chunks( + size_bytes, chunk_size + ) return response @@ -121,7 +123,9 @@ def test_image_exceeds_size_limit_without_content_length(monkeypatch): This uses the old non-streaming mock for backward compatibility. """ monkeypatch.setattr( - litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False) + litellm, + "module_level_client", + LargeImageClient(size_mb=100, include_content_length=False), ) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -134,7 +138,7 @@ def test_streaming_download_protects_against_huge_files(monkeypatch): """ Test that streaming download aborts early when file exceeds size limit, preventing memory exhaustion from huge files (e.g., petabyte-sized files). - + This test verifies that the streaming implementation doesn't download the entire file into memory before checking size. Instead, it should abort as soon as the limit is exceeded during streaming. @@ -148,7 +152,7 @@ def test_streaming_download_protects_against_huge_files(monkeypatch): # Verify the error message shows it was caught during streaming assert "exceeds maximum allowed size" in str(excinfo.value) - + # The error should be raised after downloading just slightly more than the limit # not after downloading the full 1GB @@ -187,13 +191,15 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): """ Test that streaming download can handle extremely large file URLs (e.g., petabyte-sized) without attempting to download the entire file or causing memory exhaustion. - + This simulates what happens if a malicious actor or misconfiguration provides a URL to an extremely large file. """ # Simulate a 1 petabyte file (1,000,000 GB) # Without streaming protection, this would cause OOM or hang indefinitely - client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) + client = StreamingLargeImageClient( + size_mb=1_000_000_000, include_content_length=False + ) monkeypatch.setattr(litellm, "module_level_client", client) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -214,6 +220,205 @@ def test_image_size_limit_disabled(monkeypatch): with pytest.raises(litellm.ImageFetchError) as excinfo: convert_url_to_base64("https://example.com/image.jpg") - + assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +# ============================================================================ +# Tests for Content-Type handling and URL extension inference +# ============================================================================ + + +class TestInferMediaTypeFromUrl: + """Tests for _infer_media_type_from_url function.""" + + def test_simple_png_url(self): + """Test simple PNG URL.""" + result = _infer_media_type_from_url("https://example.com/image.png") + assert result == "image/png" + + def test_simple_jpeg_url(self): + """Test simple JPEG URL with .jpeg extension.""" + result = _infer_media_type_from_url("https://example.com/photo.jpeg") + assert result == "image/jpeg" + + def test_jpg_extension(self): + """Test JPG extension maps to image/jpeg.""" + result = _infer_media_type_from_url("https://example.com/photo.jpg") + assert result == "image/jpeg" + + def test_gif_url(self): + """Test GIF URL.""" + result = _infer_media_type_from_url("https://example.com/animation.gif") + assert result == "image/gif" + + def test_webp_url(self): + """Test WebP URL.""" + result = _infer_media_type_from_url("https://example.com/image.webp") + assert result == "image/webp" + + def test_signed_url_with_query_params(self): + """Test signed URL (e.g., Aliyun OSS, AWS S3) with query parameters.""" + result = _infer_media_type_from_url( + "https://bucket.oss-cn-hangzhou.aliyuncs.com/image.png" + "?Expires=1699999999&OSSAccessKeyId=LTAI5xxx&Signature=xxx" + ) + assert result == "image/png" + + def test_signed_url_complex_query(self): + """Test signed URL with complex query string.""" + result = _infer_media_type_from_url( + "https://s3.amazonaws.com/bucket/photos/sunset.jpg" + "?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=xxx" + ) + assert result == "image/jpeg" + + def test_unsupported_extension_raises(self): + """Test that unsupported extension raises an exception.""" + with pytest.raises(Exception) as excinfo: + _infer_media_type_from_url("https://example.com/document.pdf") + assert "Unsupported image format" in str(excinfo.value) + assert "pdf" in str(excinfo.value) + + def test_case_insensitive_extension(self): + """Test that extension matching is case-insensitive.""" + result = _infer_media_type_from_url("https://example.com/IMAGE.PNG") + assert result == "image/png" + + +class TestGetValidMediaType: + """Tests for _get_valid_media_type function.""" + + def test_valid_content_type_image_png(self): + """Test valid Content-Type is returned directly.""" + result = _get_valid_media_type("image/png", "https://example.com/image.png") + assert result == "image/png" + + def test_valid_content_type_image_jpeg(self): + """Test valid JPEG Content-Type.""" + result = _get_valid_media_type("image/jpeg", "https://example.com/image.jpg") + assert result == "image/jpeg" + + def test_content_type_with_charset_parameter(self): + """Test Content-Type with charset parameter is stripped.""" + result = _get_valid_media_type( + "image/png; charset=utf-8", "https://example.com/image.png" + ) + assert result == "image/png" + + def test_content_type_with_boundary_parameter(self): + """Test Content-Type with boundary parameter is stripped.""" + result = _get_valid_media_type( + "image/jpeg; boundary=something", "https://example.com/image.jpg" + ) + assert result == "image/jpeg" + + def test_invalid_content_type_application_octet_stream(self): + """Test that application/octet-stream falls back to URL extension.""" + result = _get_valid_media_type( + "application/octet-stream", "https://example.com/image.png" + ) + assert result == "image/png" + + def test_invalid_content_type_urlencoded(self): + """ + Test that application/x-www-form-urlencoded falls back to URL extension. + This is a real-world case from Aliyun OSS CDN returning wrong Content-Type. + """ + result = _get_valid_media_type( + "application/x-www-form-urlencoded", + "https://bucket.oss-cn-hangzhou.aliyuncs.com/image.png" + "?Expires=1699999999&Signature=xxx", + ) + assert result == "image/png" + + def test_invalid_content_type_text_html(self): + """Test that text/html falls back to URL extension.""" + result = _get_valid_media_type("text/html", "https://example.com/image.webp") + assert result == "image/webp" + + def test_none_content_type(self): + """Test that None Content-Type falls back to URL extension.""" + result = _get_valid_media_type(None, "https://example.com/photo.jpeg") + assert result == "image/jpeg" + + def test_empty_content_type_after_strip(self): + """Test edge case where Content-Type is empty after stripping.""" + result = _get_valid_media_type( + "; charset=utf-8", "https://example.com/image.gif" + ) + assert result == "image/gif" + + +class InvalidContentTypeClient: + """Client that returns an invalid Content-Type.""" + + def __init__(self, content_type: str): + self.content_type = content_type + + def get(self, url, follow_redirects=True): + size_bytes = 1024 + headers = { + "Content-Type": self.content_type, + "Content-Length": str(size_bytes), + } + return Response( + status_code=200, + headers=headers, + content=b"x" * size_bytes, + request=Request("GET", url), + ) + + +def test_convert_url_handles_invalid_content_type(monkeypatch): + """ + Integration test: convert_url_to_base64 handles invalid Content-Type. + + Simulates Aliyun OSS CDN returning application/x-www-form-urlencoded for a .png file. + """ + monkeypatch.setattr( + litellm, + "module_level_client", + InvalidContentTypeClient("application/x-www-form-urlencoded"), + ) + + result = convert_url_to_base64( + "https://bucket.oss-cn-hangzhou.aliyuncs.com/image.png?Expires=xxx" + ) + + assert result.startswith("data:image/png;base64,") + + +def test_convert_url_handles_content_type_with_params(monkeypatch): + """ + Integration test: convert_url_to_base64 handles Content-Type with parameters. + + Tests that 'image/png; charset=utf-8' is properly stripped to 'image/png'. + """ + monkeypatch.setattr( + litellm, + "module_level_client", + InvalidContentTypeClient("image/png; charset=utf-8"), + ) + + result = convert_url_to_base64("https://example.com/image.png") + + assert result.startswith("data:image/png;base64,") + + +def test_convert_url_handles_application_octet_stream(monkeypatch): + """ + Integration test: convert_url_to_base64 handles application/octet-stream. + + Some servers return application/octet-stream for binary files including images. + """ + monkeypatch.setattr( + litellm, + "module_level_client", + InvalidContentTypeClient("application/octet-stream"), + ) + + result = convert_url_to_base64("https://example.com/photo.jpeg") + + assert result.startswith("data:image/jpeg;base64,") diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py index e9a1f106586..1fbfec7eed0 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py @@ -392,10 +392,10 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: """ @patch( - "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_url_to_base64" + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_to_anthropic_image_obj" ) def test_vertex_ai_messages_converts_image_url_to_base64( - self, mock_convert_url: MagicMock + self, mock_convert_obj: MagicMock ): """ Test that the /v1/messages endpoint converts image URLs to base64 for Vertex AI. @@ -403,9 +403,11 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: When using Anthropic native format with URL source type, Vertex AI should convert it to base64. """ - mock_convert_url.return_value = ( - "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ==" - ) + mock_convert_obj.return_value = { + "type": "base64", + "media_type": "image/jpeg", + "data": "/9j/4AAQSkZJRgABAQAAAQ==", + } messages = [ { @@ -426,8 +428,10 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: config = VertexAIPartnerModelsAnthropicMessagesConfig() converted = config._convert_image_urls_to_base64(messages) - # Verify convert_url_to_base64 was called - mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg") + # Verify convert_to_anthropic_image_obj was called + mock_convert_obj.assert_called_once_with( + openai_image_url="https://example.com/image.jpg", format=None + ) # Check the result has base64 source type user_message = converted[0] @@ -439,7 +443,7 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: assert image_content["source"]["data"] == "/9j/4AAQSkZJRgABAQAAAQ==" @patch( - "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_url_to_base64" + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_to_anthropic_image_obj" ) def test_vertex_ai_messages_preserves_base64_images( self, mock_convert_url: MagicMock @@ -477,17 +481,19 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: assert image_content["source"]["data"] == "iVBORw0KGgo=" @patch( - "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_url_to_base64" + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_to_anthropic_image_obj" ) def test_vertex_ai_messages_preserves_cache_control( - self, mock_convert_url: MagicMock + self, mock_convert_obj: MagicMock ): """ Test that cache_control is preserved when converting image URLs. """ - mock_convert_url.return_value = ( - "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ==" - ) + mock_convert_obj.return_value = { + "type": "base64", + "media_type": "image/jpeg", + "data": "/9j/4AAQSkZJRgABAQAAAQ==", + } messages = [ { @@ -514,7 +520,7 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: assert image_content["cache_control"] == {"type": "ephemeral"} @patch( - "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_url_to_base64" + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_to_anthropic_image_obj" ) def test_vertex_ai_messages_handles_text_only_messages( self, mock_convert_url: MagicMock @@ -542,7 +548,7 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: assert converted[0]["content"][0]["text"] == "Hello, how are you?" @patch( - "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_url_to_base64" + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_to_anthropic_image_obj" ) def test_vertex_ai_messages_handles_string_content( self, mock_convert_url: MagicMock @@ -567,17 +573,17 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: assert converted[0]["content"] == "Hello, how are you?" @patch( - "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_url_to_base64" + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation.convert_to_anthropic_image_obj" ) def test_vertex_ai_messages_converts_multiple_images( - self, mock_convert_url: MagicMock + self, mock_convert_obj: MagicMock ): """ Test that multiple image URLs in a message are all converted. """ - mock_convert_url.side_effect = [ - "data:image/jpeg;base64,/9j/image1", - "data:image/png;base64,iVBORw0image2", + mock_convert_obj.side_effect = [ + {"type": "base64", "media_type": "image/jpeg", "data": "/9j/image1"}, + {"type": "base64", "media_type": "image/png", "data": "iVBORw0image2"}, ] messages = [ @@ -607,7 +613,7 @@ class TestVertexAIAnthropicPassThroughImageURLHandling: converted = config._convert_image_urls_to_base64(messages) # Verify both URLs were converted - assert mock_convert_url.call_count == 2 + assert mock_convert_obj.call_count == 2 # Check both images are converted to base64 image1 = converted[0]["content"][1]