diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index a9ff14d6c8..019696b92f 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/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) diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 5831d7a3ec..f6832c320f 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -28,7 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( 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",