diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 75f34e3bd49..fd38bc9388d 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -3,7 +3,6 @@ Helper functions to handle images passed in messages """ import base64 -from typing import Optional from httpx import Response @@ -17,30 +16,6 @@ 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: @@ -74,18 +49,20 @@ 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: - if inferred_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: Unable to determine MIME type for url={url}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp', 'application/pdf']" + f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" ) - 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 + img_type = _img_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 9e20dc2a79b..aedc4f810cd 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://raw.githubusercontent.com/BerriAI/litellm/d769e81c90d453240c61fc572cdb27fae06a89d0/tests/llm_translation/fixtures/dummy.pdf" + "file_id": "https://cdn.jsdelivr.net/gh/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 7404845cfdf..cc13e816dde 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -232,83 +232,3 @@ 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,")