phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add test

Browse source
This commit is contained in:
duyalei 2025-04-23 18:35:41 +08:00
parent f4ae14a490
commit 74ae1f7376
2 changed files with 14 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/prompt_templates/image_handling.py
View file

@ -33,7 +33,7 @@ def _process_image_response(response: Response, url: str) -> str:
img_type = response.headers.get("Content-Type") img_type = response.headers.get("Content-Type")
if img_type not in img_type_map: if img_type not in img_type_map:
_img_type = url.split('?')[0].split(".")[-1].lower() _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: if img_type is None:
raise Exception( raise Exception(
f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']" f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']"

14
tests/llm_translation/test_prompt_factory.py
View file

@ -28,7 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.llms.vertex_ai.gemini.transformation import ( from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history, _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(): def test_llama_3_prompt():
@ -191,6 +191,18 @@ def test_convert_url_to_img():
assert "image/jpeg" in response_url 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( @pytest.mark.parametrize(
"url, expected_media_type", "url, expected_media_type",