fix _process_image_response

litellm/litellm_core_utils/prompt_templates/image_handling.py

@@ -23,24 +23,21 @@ def _process_image_response(response: Response, url: str) -> str:
 
     image_bytes = response.content
     base64_image = base64.b64encode(image_bytes).decode("utf-8")
-
-    image_type = response.headers.get("Content-Type")
-    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:
+    img_type_map = {
+        "jpg": "image/jpeg",
+        "jpeg": "image/jpeg",
+        "png": "image/png",
+        "gif": "image/gif",
+        "webp": "image/webp",
+    }
+    img_type = response.headers.get("Content-Type")
+    if img_type not in img_type_map:
+        _img_type = url.split('?')[0].split(".")[-1].lower()
+        img_type = img_type_map.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
 
     result = f"data:{img_type};base64,{base64_image}"
     in_memory_cache.set_cache(url, result)

tests/llm_translation/test_prompt_factory.py

@@ -28,7 +28,7 @@
 from litellm.llms.vertex_ai.gemini.transformation import (
     _gemini_convert_messages_with_history,
 )
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 
 def test_llama_3_prompt():
@@ -191,6 +191,18 @@ def test_convert_url_to_img():
 
     assert "image/jpeg" in response_url
 
+def test_convert_url_with_wrong_cotent_type_to_img():
+    mock_response = Mock()
+    mock_response.content = b'mock_image_data'
+    mock_response.headers = {'Content-Type': ''}
+    mock_response.status_code = 200
+    with patch('httpx.Client.get', return_value=mock_response):
+        response_url = convert_url_to_base64(
+            url="https://images.pexels.com/photos/1319515/pexels-photo-1319515.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1"
+        )
+
+    assert "image/jpeg" in response_url
+
 
 @pytest.mark.parametrize(
     "url, expected_media_type",