mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
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:
parent
f9d45f8a43
commit
cab48123d1
4 changed files with 238 additions and 17 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue