diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index a97d6812c8c4..b359145842e8 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -663,19 +663,23 @@ def convert_url_to_base64(url): image_bytes = response.content base64_image = base64.b64encode(image_bytes).decode("utf-8") - img_type = url.split(".")[-1].lower() - if img_type == "jpg" or img_type == "jpeg": - img_type = "image/jpeg" - elif img_type == "png": - img_type = "image/png" - elif img_type == "gif": - img_type = "image/gif" - elif img_type == "webp": - img_type = "image/webp" + image_type = response.headers.get("Content-Type", None) + if image_type is not None and image_type.startswith("image/"): + img_type = image_type else: - raise Exception( - f"Error: Unsupported image format. Format={img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" - ) + img_type = url.split(".")[-1].lower() + if img_type == "jpg" or img_type == "jpeg": + img_type = "image/jpeg" + elif img_type == "png": + img_type = "image/png" + elif img_type == "gif": + img_type = "image/gif" + elif img_type == "webp": + img_type = "image/webp" + else: + raise Exception( + f"Error: Unsupported image format. Format={img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" + ) return f"data:{img_type};base64,{base64_image}" else: diff --git a/litellm/tests/test_prompt_factory.py b/litellm/tests/test_prompt_factory.py index b3aafab6e6fa..5a368f92d3c2 100644 --- a/litellm/tests/test_prompt_factory.py +++ b/litellm/tests/test_prompt_factory.py @@ -1,7 +1,8 @@ #### What this tests #### # This tests if prompts are being correctly formatted -import sys import os +import sys + import pytest sys.path.insert(0, os.path.abspath("../..")) @@ -10,12 +11,13 @@ import litellm from litellm import completion from litellm.llms.prompt_templates.factory import ( - anthropic_pt, + _bedrock_tools_pt, anthropic_messages_pt, + anthropic_pt, claude_2_1_pt, + convert_url_to_base64, llama_2_chat_pt, prompt_factory, - _bedrock_tools_pt, ) @@ -153,3 +155,11 @@ def test_bedrock_tool_calling_pt(): converted_tools = _bedrock_tools_pt(tools=tools) print(converted_tools) + + +def test_convert_url_to_img(): + 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