fix(images): default null token counts in image usage on generation and edit

OpenAI-compatible image relays can return usage with null token counts:
null input_tokens / output_tokens / total_tokens, or null image_tokens /
text_tokens inside input_tokens_details. ImageUsageInputTokensDetails types
those as int, so the generation path (which only defaulted a missing
input_tokens_details) and the edit path (which had no guard) both failed
pydantic validation and the proxy returned 500 after the upstream billed.

Move the guard into normalize_image_response_usage() and use it on both
paths. Only those count fields are defaulted; output_tokens_details and
every other usage key pass through unchanged, and the edit path keeps
constructing ImageResponse through its constructor.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
pandalaohe 2026-10-15 18:17:49 +08:00
parent f9d45f8a43
commit cab48123d1
4 changed files with 238 additions and 17 deletions

View file

@ -484,6 +484,32 @@ def handle_invalid_parallel_tool_calls(
_handle_invalid_parallel_tool_calls = handle_invalid_parallel_tool_calls _handle_invalid_parallel_tool_calls = handle_invalid_parallel_tool_calls
def normalize_image_response_usage(response: Mapping[str, object]) -> Mapping[str, object]:
usage: Final = cast(Mapping[str, object] | None, response.get("usage"))
if usage is None:
return {**response}
input_tokens: Final = cast(int, usage["input_tokens"] if usage.get("input_tokens") is not None else 0)
output_tokens: Final = cast(int, usage["output_tokens"] if usage.get("output_tokens") is not None else 0)
input_details: Final = cast(Mapping[str, object], usage.get("input_tokens_details") or {})
return {
**response,
"usage": {
**usage,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": usage["total_tokens"]
if usage.get("total_tokens") is not None
else input_tokens + output_tokens,
"input_tokens_details": {
**input_details,
"image_tokens": input_details["image_tokens"] if input_details.get("image_tokens") is not None else 0,
"text_tokens": input_details["text_tokens"] if input_details.get("text_tokens") is not None else 0,
},
},
}
class LiteLLMResponseObjectHandler: class LiteLLMResponseObjectHandler:
@staticmethod @staticmethod
def convert_to_image_response( def convert_to_image_response(
@ -491,23 +517,15 @@ class LiteLLMResponseObjectHandler:
model_response_object: ImageResponse | None = None, model_response_object: ImageResponse | None = None,
hidden_params: dict | None = None, hidden_params: dict | None = None,
) -> ImageResponse: ) -> ImageResponse:
response_object.update({"hidden_params": hidden_params}) response_object.update(
{
**normalize_image_response_usage(cast(Mapping[str, object], response_object)),
"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: if "usage" in response_object and response_object["usage"] is not None:
usage: Final = response_object["usage"] usage: Final = 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,
}
# Map Responses API naming to Chat Completions API naming for cost calculator # Map Responses API naming to Chat Completions API naming for cost calculator
if usage.get("prompt_tokens") is None: if usage.get("prompt_tokens") is None:

View file

@ -6,6 +6,9 @@ from httpx._types import RequestFiles
import litellm import litellm
from litellm.images.utils import ImageEditRequestUtils from litellm.images.utils import ImageEditRequestUtils
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
normalize_image_response_usage,
)
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ( from litellm.types.images.main import (
@ -150,7 +153,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
raw_response_json: Final = raw_response.json() raw_response_json: Final = raw_response.json()
except Exception: except Exception:
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
return ImageResponse(**raw_response_json) return ImageResponse(**cast(dict[str, Any], normalize_image_response_usage(raw_response_json)))
def validate_environment( def validate_environment(
self, self,

View file

@ -1,6 +1,7 @@
import asyncio import asyncio
import importlib import importlib
import json import json
from copy import deepcopy
from datetime import datetime, timedelta from datetime import datetime, timedelta
from typing import Final from typing import Final
@ -8,7 +9,15 @@ import pytest
import litellm import litellm
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import LiteLLMResponseObjectHandler, convert_to_model_response_object, convert_to_streaming_response, handle_invalid_parallel_tool_calls, safe_convert_created_field, should_convert_tool_call_to_json_mode from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
LiteLLMResponseObjectHandler,
convert_to_model_response_object,
convert_to_streaming_response,
handle_invalid_parallel_tool_calls,
normalize_image_response_usage,
safe_convert_created_field,
should_convert_tool_call_to_json_mode,
)
from litellm.types.utils import ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, Choices, CompletionTokensDetailsWrapper, Function, ImageObject, ImageResponse, Message, ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.types.utils import ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, Choices, CompletionTokensDetailsWrapper, Function, ImageObject, ImageResponse, Message, ModelResponse, PromptTokensDetailsWrapper, Usage
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -373,6 +382,97 @@ def test_convert_to_image_response_with_extra_fields_2():
assert result.data[0].url == "http://example.com/image1.jpg" assert result.data[0].url == "http://example.com/image1.jpg"
assert result.data[1].url == "http://example.com/image2.jpg" assert result.data[1].url == "http://example.com/image2.jpg"
@pytest.mark.parametrize("response", [{}, {"usage": None}])
def test_normalize_image_response_without_usage_returns_copy(response: dict[str, object]) -> None:
normalized: Final = normalize_image_response_usage(response)
assert normalized == response
assert normalized is not response
@pytest.mark.parametrize(
("usage_fields", "expected_fields"),
[
pytest.param(
{"input_tokens_details": {"image_tokens": None, "text_tokens": 5}},
{"input_tokens_details": {"image_tokens": 0, "text_tokens": 5}},
id="null-input-image-tokens",
),
pytest.param(
{"input_tokens_details": {"image_tokens": 6}},
{"input_tokens_details": {"image_tokens": 6, "text_tokens": 0}},
id="missing-input-text-tokens",
),
pytest.param(
{"input_tokens_details": {"text_tokens": 5}},
{"input_tokens_details": {"image_tokens": 0, "text_tokens": 5}},
id="missing-input-image-tokens",
),
pytest.param(
{"output_tokens_details": None},
{"output_tokens_details": None},
id="null-output-details",
),
pytest.param(
{"input_tokens_details": {"image_tokens": 6, "text_tokens": 4, "cached_tokens_details": None}},
{"input_tokens_details": {"image_tokens": 6, "text_tokens": 4, "cached_tokens_details": None}},
id="null-cached-details",
),
pytest.param(
{"output_tokens_details": {"image_tokens": 20}},
{"output_tokens_details": {"image_tokens": 20}},
id="sparse-output-details",
),
pytest.param(
{"output_tokens_details": {"image_tokens": None}},
{"output_tokens_details": {"image_tokens": None}},
id="null-output-image-tokens",
),
pytest.param(
{"output_tokens_details": {"image_tokens": 18, "text_tokens": 2}},
{"output_tokens_details": {"image_tokens": 18, "text_tokens": 2}},
id="fully-populated-usage",
),
pytest.param(
{"provider_metadata": None},
{"provider_metadata": None},
id="null-provider-metadata",
),
pytest.param({}, {}, id="absent-output-details"),
],
)
def test_convert_to_image_response_preserves_usage_details(
usage_fields: dict[str, object], expected_fields: dict[str, object]
) -> None:
usage: Final = {
"input_tokens": 10,
"input_tokens_details": {"image_tokens": 6, "text_tokens": 4},
"output_tokens": 20,
"total_tokens": 30,
}
response: Final = {
"created": 1234567890,
"data": [{"b64_json": "generated-image"}],
"usage": {**usage, **usage_fields},
}
original: Final = deepcopy(response)
expected_usage: Final = {**usage, **expected_fields}
normalized: Final = normalize_image_response_usage(response)
assert normalized == {**response, "usage": expected_usage}
assert normalized is not response
assert response == original
result: Final = LiteLLMResponseObjectHandler.convert_to_image_response(response)
assert result.usage is not None
serialized_usage: Final = result.usage.model_dump()
assert {key: serialized_usage[key] for key in expected_usage} == expected_usage
assert result.data[0].b64_json == "generated-image"
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_convert_to_image_response_with_none_usage_fields(): def test_convert_to_image_response_with_none_usage_fields():
""" """

View file

@ -1,12 +1,15 @@
from io import BufferedReader, BytesIO from io import BufferedReader, BytesIO
from typing import Dict from typing import Dict, Final
from unittest.mock import Mock, patch
import httpx
import pytest import pytest
from litellm import image_edit from litellm import image_edit
from litellm.images.utils import ImageEditRequestUtils from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.types.router import GenericLiteLLMParams from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageObject
@pytest.fixture @pytest.fixture
@ -14,6 +17,103 @@ def image_edit_config() -> OpenAIImageEditConfig:
return OpenAIImageEditConfig() return OpenAIImageEditConfig()
@pytest.mark.parametrize(
("usage_fields", "expected_fields"),
[
pytest.param(
{"input_tokens": None, "output_tokens": None, "total_tokens": None},
{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
id="null-counts",
),
pytest.param(
{"input_tokens_details": None},
{"input_tokens_details": {"image_tokens": 0, "text_tokens": 0}},
id="null-input-details",
),
pytest.param(
{"input_tokens_details": {"image_tokens": None, "text_tokens": None}},
{"input_tokens_details": {"image_tokens": 0, "text_tokens": 0}},
id="null-input-detail-counts",
),
pytest.param(
{"input_tokens_details": {"image_tokens": None, "text_tokens": 5}},
{"input_tokens_details": {"image_tokens": 0, "text_tokens": 5}},
id="mixed-input-detail-counts",
),
pytest.param(
{"output_tokens_details": {"image_tokens": None, "text_tokens": None}},
{"output_tokens_details": {"image_tokens": None, "text_tokens": None}},
id="null-output-detail-counts",
),
pytest.param(
{"output_tokens_details": {"image_tokens": 18, "text_tokens": None}},
{"output_tokens_details": {"image_tokens": 18, "text_tokens": None}},
id="mixed-output-detail-counts",
),
pytest.param(
{"output_tokens_details": None},
{"output_tokens_details": None},
id="null-output-details",
),
pytest.param(
{"input_tokens_details": {"image_tokens": 6, "text_tokens": 4, "cached_tokens_details": None}},
{"input_tokens_details": {"image_tokens": 6, "text_tokens": 4, "cached_tokens_details": None}},
id="null-cached-details",
),
pytest.param(
{"output_tokens_details": {"image_tokens": 20}},
{"output_tokens_details": {"image_tokens": 20}},
id="sparse-output-details",
),
pytest.param({"total_tokens": None}, {"total_tokens": 30}, id="calculated-total"),
pytest.param({}, {}, id="integer-usage-unchanged"),
],
)
def test_transform_image_edit_response_normalizes_null_usage(
image_edit_config: OpenAIImageEditConfig,
usage_fields: dict[str, object],
expected_fields: dict[str, object],
) -> None:
usage: Final = {
"input_tokens": 10,
"output_tokens": 20,
"total_tokens": 30,
"input_tokens_details": {"image_tokens": 6, "text_tokens": 4},
"output_tokens_details": {"image_tokens": 18, "text_tokens": 2},
}
raw_response: Final = httpx.Response(
status_code=200,
json={"created": 1234567890, "data": [{"b64_json": "edited-image"}], "usage": {**usage, **usage_fields}},
)
result: Final = image_edit_config.transform_image_edit_response(
model="gpt-image-1", raw_response=raw_response, logging_obj=Mock()
)
assert result.usage is not None
assert result.usage.model_dump() == {**usage, **expected_fields}
assert result.data[0].b64_json == "edited-image"
def test_transform_image_edit_response_defaults_created_and_converts_image_objects(
image_edit_config: OpenAIImageEditConfig,
) -> None:
raw_response: Final = httpx.Response(
status_code=200,
json={"data": [{"b64_json": "edited-image"}]},
)
with patch("time.time", return_value=1234567890):
result: Final = image_edit_config.transform_image_edit_response(
model="gpt-image-1", raw_response=raw_response, logging_obj=Mock()
)
assert result.created == 1234567890
assert len(result.data) == 1
assert isinstance(result.data[0], ImageObject)
assert result.data[0].b64_json == "edited-image"
def test_transform_image_edit_request_basic(image_edit_config: OpenAIImageEditConfig): def test_transform_image_edit_request_basic(image_edit_config: OpenAIImageEditConfig):
"""Test basic image edit request transformation with image and prompt""" """Test basic image edit request transformation with image and prompt"""
model = "dall-e-2" model = "dall-e-2"