diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2d0e262fe..942750416 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1145,7 +1145,7 @@ class ImageObject(OpenAIImage): url: Optional[str] = None revised_prompt: Optional[str] = None - def __init__(self, b64_json=None, url=None, revised_prompt=None): + def __init__(self, b64_json=None, url=None, revised_prompt=None, **kwargs): super().__init__(b64_json=b64_json, url=url, revised_prompt=revised_prompt) # type: ignore def __contains__(self, key): diff --git a/tests/llm_translation/test_convert_dict_to_image.py b/tests/llm_translation/test_convert_dict_to_image.py new file mode 100644 index 000000000..87c415ecb --- /dev/null +++ b/tests/llm_translation/test_convert_dict_to_image.py @@ -0,0 +1,119 @@ +import json +import os +import sys +from datetime import datetime + +sys.path.insert( + 0, os.path.abspath("../../") +) # Adds the parent directory to the system path + +import litellm +import pytest +from datetime import timedelta +from litellm.types.utils import ImageResponse, ImageObject +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler, +) + + +def test_convert_to_image_response_basic(): + # Test basic conversion with minimal input + response_dict = { + "created": 1234567890, + "data": [{"url": "http://example.com/image.jpg"}], + } + + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert isinstance(result, ImageResponse) + assert result.created == 1234567890 + assert result.data[0].url == "http://example.com/image.jpg" + + +def test_convert_to_image_response_with_hidden_params(): + # Test with hidden params + response_dict = { + "created": 1234567890, + "data": [{"url": "http://example.com/image.jpg"}], + } + hidden_params = {"api_key": "test_key"} + + result = LiteLLMResponseObjectHandler.convert_to_image_response( + response_dict, hidden_params=hidden_params + ) + + assert result._hidden_params == {"api_key": "test_key"} + + +def test_convert_to_image_response_multiple_images(): + # Test handling multiple images in response + response_dict = { + "created": 1234567890, + "data": [ + {"url": "http://example.com/image1.jpg"}, + {"url": "http://example.com/image2.jpg"}, + ], + } + + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert len(result.data) == 2 + assert result.data[0].url == "http://example.com/image1.jpg" + assert result.data[1].url == "http://example.com/image2.jpg" + + +def test_convert_to_image_response_with_b64_json(): + # Test handling b64_json in response + response_dict = { + "created": 1234567890, + "data": [{"b64_json": "base64encodedstring"}], + } + + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert result.data[0].b64_json == "base64encodedstring" + + +def test_convert_to_image_response_with_extra_fields(): + response_dict = { + "created": 1234567890, + "data": [ + { + "url": "http://example.com/image1.jpg", + "content_filter_results": {"category": "violence", "flagged": True}, + }, + { + "url": "http://example.com/image2.jpg", + "content_filter_results": {"category": "violence", "flagged": True}, + }, + ], + } + + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert result.data[0].url == "http://example.com/image1.jpg" + assert result.data[1].url == "http://example.com/image2.jpg" + + +def test_convert_to_image_response_with_extra_fields_2(): + """ + Date from a non-OpenAI API could have some obscure field in addition to the expected ones. This should not break the conversion. + """ + response_dict = { + "created": 1234567890, + "data": [ + { + "url": "http://example.com/image1.jpg", + "very_obscure_field": "some_value", + }, + { + "url": "http://example.com/image2.jpg", + "very_obscure_field2": "some_other_value", + }, + ], + } + + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert result.data[0].url == "http://example.com/image1.jpg" + assert result.data[1].url == "http://example.com/image2.jpg"