diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 63bbd6dc904..ba6c41eef45 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -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: diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index d54522597a0..36917532045 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -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, diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index f918d4cb79e..0f4bc086a94 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -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(): """ diff --git a/tests/unit/llms/openai/test_openai_image_edit_transformation.py b/tests/unit/llms/openai/test_openai_image_edit_transformation.py index da3352a4943..b5b03cf825f 100644 --- a/tests/unit/llms/openai/test_openai_image_edit_transformation.py +++ b/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"