From 7c513856dcf40a5b858979a0c818a8cfd1f784df Mon Sep 17 00:00:00 2001 From: Cole McIntosh <82463175+colesmcintosh@users.noreply.github.com> Date: Thu, 5 Jun 2025 14:19:18 -0600 Subject: [PATCH] Fix None values in usage field for gpt-image-1 model responses (#11448) * fix(convert_dict_to_response.py): handle None values in usage field for gpt-image-1 * test: add tests for handling None and partial values in usage fields for gpt-image-1 responses --- .../convert_dict_to_response.py | 16 +++ .../test_convert_dict_to_image.py | 104 ++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 26134153511..a44d6c29e01 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -294,6 +294,22 @@ class LiteLLMResponseObjectHandler: ) -> ImageResponse: response_object.update({"hidden_params": hidden_params}) + # Handle gpt-image-1 usage field with None values + if "usage" in response_object and response_object["usage"] is not None: + usage = response_object["usage"] + # Check if usage fields are None and provide defaults + if usage.get("input_tokens") is None: + usage["input_tokens"] = 0 + if usage.get("output_tokens") is None: + usage["output_tokens"] = 0 + if usage.get("total_tokens") is None: + usage["total_tokens"] = usage["input_tokens"] + usage["output_tokens"] + if usage.get("input_tokens_details") is None: + usage["input_tokens_details"] = { + "image_tokens": 0, + "text_tokens": 0, + } + if model_response_object is None: model_response_object = ImageResponse(**response_object) return model_response_object diff --git a/tests/llm_translation/test_convert_dict_to_image.py b/tests/llm_translation/test_convert_dict_to_image.py index 87c415ecb31..d82b8deaf09 100644 --- a/tests/llm_translation/test_convert_dict_to_image.py +++ b/tests/llm_translation/test_convert_dict_to_image.py @@ -117,3 +117,107 @@ def test_convert_to_image_response_with_extra_fields_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_none_usage_fields(): + """ + Test handling of None values in usage fields, specifically for gpt-image-1 responses. + + This test verifies the fix for the bug where gpt-image-1 returns None values + for usage statistics fields, which caused Pydantic validation errors. + The fix should clean these None values and let ImageResponse constructor + handle the default values. + """ + response_dict = { + "created": 1234567890, + "data": [{"b64_json": "base64encodedstring"}], + "usage": { + "input_tokens": None, # gpt-image-1 returns None instead of integer + "input_tokens_details": None, # gpt-image-1 returns None instead of object + "output_tokens": None, # gpt-image-1 returns None instead of integer + "total_tokens": None, # gpt-image-1 returns None instead of integer + } + } + + # This should not raise a ValidationError + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert isinstance(result, ImageResponse) + assert result.created == 1234567890 + assert result.data[0].b64_json == "base64encodedstring" + + # Usage should be properly initialized with default values + assert result.usage is not None + assert result.usage.input_tokens == 0 + assert result.usage.output_tokens == 0 + assert result.usage.total_tokens == 0 + assert result.usage.input_tokens_details is not None + assert result.usage.input_tokens_details.image_tokens == 0 + assert result.usage.input_tokens_details.text_tokens == 0 + + +def test_convert_to_image_response_with_partial_none_usage_fields(): + """ + Test handling of mixed None and valid values in usage fields. + """ + response_dict = { + "created": 1234567890, + "data": [{"b64_json": "base64encodedstring"}], + "usage": { + "input_tokens": 10, # Valid value + "input_tokens_details": None, # None value (should be cleaned) + "output_tokens": None, # None value (should be cleaned) + "total_tokens": 10, # Valid value + } + } + + # This should not raise a ValidationError + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert isinstance(result, ImageResponse) + assert result.created == 1234567890 + assert result.data[0].b64_json == "base64encodedstring" + + # Usage should be properly initialized with defaults where needed + # Valid values should be preserved, None values should be cleaned and use defaults + assert result.usage is not None + assert result.usage.input_tokens == 10 # Valid value should be preserved + assert result.usage.output_tokens == 0 # None value should become 0 + assert result.usage.total_tokens == 10 # Calculated as input_tokens + output_tokens (10 + 0) + assert result.usage.input_tokens_details is not None + assert result.usage.input_tokens_details.image_tokens == 0 + assert result.usage.input_tokens_details.text_tokens == 0 + + +def test_convert_to_image_response_with_valid_usage_fields(): + """ + Test that valid usage fields are preserved correctly. + """ + response_dict = { + "created": 1234567890, + "data": [{"b64_json": "base64encodedstring"}], + "usage": { + "input_tokens": 50, + "input_tokens_details": { + "image_tokens": 30, + "text_tokens": 20, + }, + "output_tokens": 10, + "total_tokens": 60, + } + } + + result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict) + + assert isinstance(result, ImageResponse) + assert result.created == 1234567890 + assert result.data[0].b64_json == "base64encodedstring" + + # Valid usage fields should be preserved + assert result.usage is not None + assert result.usage.input_tokens == 50 + assert result.usage.output_tokens == 10 + assert result.usage.total_tokens == 60 + assert result.usage.input_tokens_details is not None + assert result.usage.input_tokens_details.image_tokens == 30 + assert result.usage.input_tokens_details.text_tokens == 20