diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index bcdfc70c1a..019696b92f 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -33,7 +33,7 @@ def _process_image_response(response: Response, url: str) -> str: 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) + 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']" 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",