From f810a4c0367eff2aa9b0d13c2a54dacbf74b55cf Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 21 May 2025 21:30:34 -0700 Subject: [PATCH] fix: test --- litellm/types/utils.py | 9 ++------- tests/image_gen_tests/base_image_generation_test.py | 4 +++- 2 files changed, 5 insertions(+), 8 deletions(-) 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)