fix(image-handling): infer the image mime type when the server sends a generic content type

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/litellm_core_utils/prompt_templates/image_handling.py

@@ -15,6 +15,7 @@ import litellm
 from litellm import verbose_logger
 from litellm.caching.caching import InMemoryCache
 from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
+from litellm.litellm_core_utils.prompt_templates.common_utils import infer_content_type_from_url_and_content
 from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
 from litellm.types.llms.openai import AllMessageValues
 
@@ -55,23 +56,16 @@ def _process_image_response(response: Response, url: str) -> str:
 
     base64_image: Final = base64.b64encode(image_bytes).decode("utf-8")
 
-    image_type: Final = response.headers.get("Content-Type")
-    if image_type is None:
-        img_type = url.split(".")[-1].lower()
-        _img_type: Final = {
-            "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
+    try:
+        img_type: Final = infer_content_type_from_url_and_content(
+            url=url,
+            content=bytes(image_bytes),
+            current_content_type=response.headers.get("Content-Type"),
+        )
+    except ValueError as e:
+        raise litellm.ImageFetchError(
+            f"Error: Unable to determine image content type from the server's headers, the URL, or the image bytes. url={url}"
+        ) from e
 
     result: Final = f"data:{img_type};base64,{base64_image}"
     in_memory_cache.set_cache(url, result)

tests/unit/litellm_core_utils/test_image_handling.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import copy
 import time
 import uuid
@@ -259,6 +260,54 @@ async def test_async_data_url_is_returned_unchanged_without_fetch(monkeypatch):
     assert await async_convert_url_to_base64(data_url) == data_url
 
 
+REAL_PNG_BYTES = base64.b64decode(
+    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
+)
+
+
+def _stub_image_client(content, content_type):
+    class _Client:
+        def get(self, url, follow_redirects=True):
+            headers = {} if content_type is None else {"Content-Type": content_type}
+            return Response(200, content=content, headers=headers, request=Request("GET", url))
+
+    return _Client()
+
+
+def test_convert_url_to_base64_infers_the_type_when_the_server_sends_octet_stream(monkeypatch):
+    monkeypatch.setattr(
+        litellm, "module_level_client", _stub_image_client(REAL_PNG_BYTES, "application/octet-stream")
+    )
+
+    result = convert_url_to_base64(f"http://img.example/{uuid.uuid4()}")
+
+    assert result.startswith("data:image/png;base64,")
+
+
+def test_convert_url_to_base64_keeps_a_real_content_type(monkeypatch):
+    monkeypatch.setattr(
+        litellm, "module_level_client", _stub_image_client(REAL_PNG_BYTES, "image/jpeg")
+    )
+
+    result = convert_url_to_base64(f"http://img.example/{uuid.uuid4()}.png")
+
+    assert result.startswith("data:image/jpeg;base64,")
+
+
+def test_convert_url_to_base64_raises_when_no_content_type_is_determinable(monkeypatch):
+    monkeypatch.setattr(
+        litellm,
+        "module_level_client",
+        _stub_image_client(b"\x00\x01\x02\x03not-an-image", "application/octet-stream"),
+    )
+    url = f"http://img.example/{uuid.uuid4()}"
+
+    with pytest.raises(litellm.ImageFetchError) as excinfo:
+        convert_url_to_base64(url)
+
+    assert url in str(excinfo.value)
+
+
 def test_image_size_limit_disabled(monkeypatch):
     """
     Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables all image URL downloads.