mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +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
|
_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:
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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():
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue