diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index cb5c02ce10e..57b4f545301 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/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) diff --git a/tests/unit/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py index 1fc8d54cecc..57eb32f98e5 100644 --- a/tests/unit/litellm_core_utils/test_image_handling.py +++ b/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.