diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 89a766de05c..5d4b20bdb0d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1567,7 +1567,7 @@ class ImageUsage(BaseLiteLLMOpenAIResponseObject): from openai.types.images_response import ImagesResponse as OpenAIImageResponse -class ImageResponse(BaseLiteLLMOpenAIResponseObject): +class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = {} usage: ImageUsage @@ -1609,12 +1609,7 @@ class ImageResponse(BaseLiteLLMOpenAIResponseObject): output_tokens=0, total_tokens=0, ) - - ####################### - self.created = created - self.data = _data - self.usage = _usage - ####################### + super().__init__(created=created, data=_data, usage=_usage) # type: ignore self._hidden_params = hidden_params or {} def __contains__(self, key): diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py index 7ffb2b59dea..98c00f305eb 100644 --- a/tests/image_gen_tests/base_image_generation_test.py +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -71,7 +71,9 @@ class BaseImageGenTest(ABC): from openai.types.images_response import ImagesResponse - ImagesResponse.model_validate(dict(response)) + response_dict = response.model_dump() + print("response usage=", response_dict.get("usage")) + ImagesResponse.model_validate(response_dict) for d in response.data: assert isinstance(d, Image)