From bd2eb687a235be270f6f9a2a4a2b3a4beffdb823 Mon Sep 17 00:00:00 2001 From: Jerry-Xin Date: Sat, 7 Mar 2026 18:14:40 +0800 Subject: [PATCH] fix(vertex_ai): raise ImageFetchError for unsupported image formats Change _infer_media_type_from_url to raise litellm.ImageFetchError instead of generic Exception. This ensures unsupported format errors are properly propagated through convert_url_to_base64 functions instead of being silently swallowed. --- .../prompt_templates/image_handling.py | 2 +- .../litellm_core_utils/test_image_handling.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 29db1df5832..a2337214e59 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -43,7 +43,7 @@ def _infer_media_type_from_url(url: str) -> str: extension = last_segment.rsplit(".", 1)[-1].lower() if "." in last_segment else "" media_type = EXTENSION_TO_MEDIA_TYPE.get(extension) if media_type is None: - raise Exception( + raise litellm.ImageFetchError( f"Error: Unsupported image format. Could not infer media type from URL '{url}'. " f"Supported types = {list(SUPPORTED_IMAGE_TYPES)}" ) 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 961f660cb1c..6f5167eaf57 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -275,8 +275,8 @@ class TestInferMediaTypeFromUrl: assert result == "image/jpeg" def test_unsupported_extension_raises(self): - """Test that unsupported extension raises an exception.""" - with pytest.raises(Exception) as excinfo: + """Test that unsupported extension raises ImageFetchError.""" + with pytest.raises(litellm.ImageFetchError) 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) @@ -287,9 +287,9 @@ class TestInferMediaTypeFromUrl: assert result == "image/png" def test_url_without_extension_raises_with_clear_message(self): - """Test that URL without extension raises an exception with a clear message.""" + """Test that URL without extension raises ImageFetchError with a clear message.""" url = "https://cdn.example.com/images/abc123" - with pytest.raises(Exception) as excinfo: + with pytest.raises(litellm.ImageFetchError) as excinfo: _infer_media_type_from_url(url) error_msg = str(excinfo.value) assert "Unsupported image format" in error_msg @@ -298,18 +298,18 @@ class TestInferMediaTypeFromUrl: assert "Supported types" in error_msg def test_url_with_trailing_slash_no_extension(self): - """Test URL with trailing slash and no extension.""" + """Test URL with trailing slash and no extension raises ImageFetchError.""" url = "https://cdn.example.com/images/abc123/" - with pytest.raises(Exception) as excinfo: + with pytest.raises(litellm.ImageFetchError) as excinfo: _infer_media_type_from_url(url) error_msg = str(excinfo.value) assert "Unsupported image format" in error_msg assert url in error_msg def test_url_with_dots_in_path_but_no_image_extension(self): - """Test URL with dots in path segments but no valid image extension.""" + """Test URL with dots in path segments but no valid image extension raises ImageFetchError.""" url = "https://api.example.com/v1.0/images/get" - with pytest.raises(Exception) as excinfo: + with pytest.raises(litellm.ImageFetchError) as excinfo: _infer_media_type_from_url(url) error_msg = str(excinfo.value) assert "Unsupported image format" in error_msg