diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index fd38bc9388..75f34e3bd4 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -3,6 +3,7 @@ Helper functions to handle images passed in messages """ import base64 +from typing import Optional from httpx import Response @@ -16,6 +17,30 @@ MAX_IMGS_IN_MEMORY = 10 in_memory_cache = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) +_URL_EXTENSION_TO_MIME_TYPE = { + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "png": "image/png", + "gif": "image/gif", + "webp": "image/webp", + "pdf": "application/pdf", + "txt": "text/plain", +} + +_GENERIC_BINARY_CONTENT_TYPES = ("application/octet-stream", "binary/octet-stream") + + +def _infer_mime_type_from_url(url: str) -> Optional[str]: + path = url.split("?", 1)[0].split("#", 1)[0] + ext = path.rsplit(".", 1)[-1].lower() if "." in path else "" + return _URL_EXTENSION_TO_MIME_TYPE.get(ext) + + +def _is_generic_binary_content_type(content_type: str) -> bool: + return ( + content_type.split(";", 1)[0].strip().lower() in _GENERIC_BINARY_CONTENT_TYPES + ) + def _process_image_response(response: Response, url: str) -> str: if response.status_code != 200: @@ -49,20 +74,18 @@ def _process_image_response(response: Response, url: str) -> str: base64_image = base64.b64encode(image_bytes).decode("utf-8") image_type = response.headers.get("Content-Type") + inferred_type = _infer_mime_type_from_url(url) 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: + if inferred_type is None: raise Exception( - f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" + f"Error: Unable to determine MIME type for url={url}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp', 'application/pdf']" ) - img_type = _img_type + img_type = inferred_type + elif _is_generic_binary_content_type(image_type) and inferred_type is not None: + # Some hosts (e.g. raw.githubusercontent.com, GitHub releases) serve PDFs and + # other binaries as application/octet-stream. Trust the URL extension when the + # response Content-Type carries no useful signal. + img_type = inferred_type else: img_type = image_type diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index aedc4f810c..9e20dc2a79 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -401,7 +401,7 @@ class BaseLLMChatTest(ABC): { "type": "file", "file": { - "file_id": "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/tests/llm_translation/fixtures/dummy.pdf" + "file_id": "https://raw.githubusercontent.com/BerriAI/litellm/d769e81c90d453240c61fc572cdb27fae06a89d0/tests/llm_translation/fixtures/dummy.pdf" }, }, ] 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 cc13e816dd..7404845cfd 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -232,3 +232,83 @@ def test_image_size_limit_disabled(monkeypatch): assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +class _FixedContentTypeClient: + def __init__(self, content_type=None): + self._content_type = content_type + + def get(self, url, follow_redirects=True): + body = b"x" * 1024 + headers = {"Content-Length": str(len(body))} + if self._content_type is not None: + headers["Content-Type"] = self._content_type + return Response( + status_code=200, + headers=headers, + content=body, + request=Request("GET", url), + ) + + +def test_octet_stream_pdf_falls_back_to_url_extension(monkeypatch): + """ + raw.githubusercontent.com and GH releases serve PDFs as application/octet-stream. + The fetcher must fall back to the URL extension so OpenAI/Gemini accept the data URI. + """ + monkeypatch.setattr( + litellm, + "module_level_client", + _FixedContentTypeClient(content_type="application/octet-stream"), + ) + + result = convert_url_to_base64( + "https://raw.githubusercontent.com/example/repo/main/doc.pdf" + ) + + assert result.startswith("data:application/pdf;base64,") + + +def test_octet_stream_with_unknown_extension_passes_through(monkeypatch): + """ + If the URL extension is not recognized, fall through to the server-provided + Content-Type rather than guessing. Don't break callers that do tolerate + application/octet-stream. + """ + monkeypatch.setattr( + litellm, + "module_level_client", + _FixedContentTypeClient(content_type="application/octet-stream"), + ) + + result = convert_url_to_base64("https://example.com/file.bin") + + assert result.startswith("data:application/octet-stream;base64,") + + +def test_specific_content_type_is_trusted_over_extension(monkeypatch): + """ + A meaningful Content-Type (e.g. image/png) must beat the URL extension — + we only override on generic binary types. + """ + monkeypatch.setattr( + litellm, + "module_level_client", + _FixedContentTypeClient(content_type="image/png"), + ) + + result = convert_url_to_base64("https://example.com/photo.pdf") + + assert result.startswith("data:image/png;base64,") + + +def test_missing_content_type_uses_extension(monkeypatch): + monkeypatch.setattr( + litellm, + "module_level_client", + _FixedContentTypeClient(content_type=None), + ) + + result = convert_url_to_base64("https://example.com/photo.png") + + assert result.startswith("data:image/png;base64,")