Merge 607d3fff6f into 1c6d55c7e8

litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py

@@ -530,6 +530,37 @@ 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]:
+    raw_usage: Final = response.get("usage")
+    if not isinstance(raw_usage, Mapping):
+        return {**response}
+
+    usage: Final = _FIELD_MAP.validate_python(raw_usage)
+    input_tokens: Final = usage["input_tokens"] if usage.get("input_tokens") is not None else 0
+    output_tokens: Final = usage["output_tokens"] if usage.get("output_tokens") is not None else 0
+    total_tokens: Final = (
+        input_tokens + output_tokens if isinstance(input_tokens, int) and isinstance(output_tokens, int) else None
+    )
+    raw_input_details: Final = usage.get("input_tokens_details")
+    input_details: Final = (
+        _FIELD_MAP.validate_python(raw_input_details) if isinstance(raw_input_details, Mapping) else {}
+    )
+    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 total_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(
@@ -537,23 +568,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(_FIELD_MAP.validate_python(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:

litellm/llms/openai/image_edit/transformation.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from io import BufferedReader
 from typing import TYPE_CHECKING, Any, Final, cast
 
@@ -6,6 +7,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 (
@@ -151,7 +155,8 @@ 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)
+        normalized_response: Final[Mapping[str, Any]] = normalize_image_response_usage(raw_response_json)
+        return ImageResponse(**normalized_response)
 
     def validate_environment(
         self,

tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py

@@ -1,6 +1,8 @@
 import asyncio
 import importlib
 import json
+from collections.abc import Mapping
+from copy import deepcopy
 from datetime import datetime, timedelta
 from typing import Final
 
@@ -8,7 +10,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
 
@@ -421,6 +431,122 @@ 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(
+    "counts",
+    [
+        {"input_tokens": "10", "output_tokens": 20},
+        {"input_tokens": 10, "output_tokens": "20"},
+        {"input_tokens": 1.5, "output_tokens": 20},
+        {"input_tokens": 10, "output_tokens": []},
+    ],
+)
+def test_normalize_image_response_preserves_non_integer_counts(counts: Mapping[str, object]) -> None:
+    response: Final = {"usage": {**counts, "total_tokens": None}}
+    original: Final = deepcopy(response)
+
+    normalized: Final = normalize_image_response_usage(response)
+
+    assert normalized == {
+        "usage": {
+            **counts,
+            "total_tokens": None,
+            "input_tokens_details": {"image_tokens": 0, "text_tokens": 0},
+        }
+    }
+    assert response == original
+
+
+@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():
     """

tests/unit/llms/openai/test_openai_image_edit_transformation.py

@@ -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"