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
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:
@staticmethod
def convert_to_image_response(
@ -491,23 +517,15 @@ class LiteLLMResponseObjectHandler:
model_response_object: ImageResponse | None = None,
hidden_params: dict | None = None,
) -> 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:
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
if usage.get("prompt_tokens") is None:

View file

@ -6,6 +6,9 @@ from httpx._types import RequestFiles
import litellm
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.secret_managers.main import get_secret_str
from litellm.types.images.main import (
@ -150,7 +153,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
raw_response_json: Final = raw_response.json()
except Exception:
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(
self,

View file

@ -1,6 +1,7 @@
import asyncio
import importlib
import json
from copy import deepcopy
from datetime import datetime, timedelta
from typing import Final
@ -8,7 +9,15 @@ import pytest
import litellm
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 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[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")
def test_convert_to_image_response_with_none_usage_fields():
"""

View file

@ -1,12 +1,15 @@
from io import BufferedReader, BytesIO
from typing import Dict
from typing import Dict, Final
from unittest.mock import Mock, patch
import httpx
import pytest
from litellm import image_edit
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageObject
@pytest.fixture
@ -14,6 +17,103 @@ def image_edit_config() -> 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):
"""Test basic image edit request transformation with image and prompt"""
model = "dall-e-2"