From a76d9919521e0723b412e222fa0359d91f80f994 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 27 May 2025 17:52:15 -0700 Subject: [PATCH] [Fix]: Add cost tracking for image edits endpoint [OpenAI, Azure] (#11186) * fix: add cost tracking for image edits * fix: add cost tracking for azure image edits * fix: fix linting error * fix: fixes for background param * fix: image edit test fixes * fix: openai image edit cost tracking --- litellm/cost_calculator.py | 34 ++-- .../litellm_core_utils/llm_cost_calc/utils.py | 27 ++- litellm/types/utils.py | 7 + tests/image_gen_tests/test_image_edits.py | 170 ++++++++++++++++++ 4 files changed, 221 insertions(+), 17 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 041e8b4c388..be2190adda3 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + CostCalculatorUtils, _generic_cost_per_character, generic_cost_per_token, select_cost_metric_for_model, @@ -73,7 +74,6 @@ from litellm.types.utils import ( LlmProviders, LlmProvidersSet, ModelInfo, - PassthroughCallTypes, StandardBuiltInToolsParams, Usage, ) @@ -746,12 +746,7 @@ def completion_cost( # noqa: PLR0915 str(e) ) ) - if ( - call_type == CallTypes.image_generation.value - or call_type == CallTypes.aimage_generation.value - or call_type - == PassthroughCallTypes.passthrough_image_generation.value - ): + if CostCalculatorUtils._call_type_has_image_response(call_type): ### IMAGE GENERATION COST CALCULATION ### if custom_llm_provider == "vertex_ai": if isinstance(completion_response, ImageResponse): @@ -1114,9 +1109,13 @@ def default_image_cost_calculator( # Build model names for cost lookup base_model_name = f"{size_str}/{model}" - if custom_llm_provider and model.startswith(custom_llm_provider): + model_name_without_custom_llm_provider: Optional[str] = None + if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): + model_name_without_custom_llm_provider = model.replace( + f"{custom_llm_provider}/", "" + ) base_model_name = ( - f"{custom_llm_provider}/{size_str}/{model.replace(custom_llm_provider, '')}" + f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}" ) model_name_with_quality = ( f"{quality}/{base_model_name}" if quality else base_model_name @@ -1138,17 +1137,18 @@ def default_image_cost_calculator( # Try model with quality first, fall back to base model name cost_info: Optional[dict] = None - models_to_check = [ + models_to_check: List[Optional[str]] = [ model_name_with_quality, base_model_name, model_name_with_v2_quality, model_with_quality_without_provider, model_without_provider, model, + model_name_without_custom_llm_provider, ] - for model in models_to_check: - if model in litellm.model_cost: - cost_info = litellm.model_cost[model] + for _model in models_to_check: + if _model is not None and _model in litellm.model_cost: + cost_info = litellm.model_cost[_model] break if cost_info is None: raise Exception( @@ -1353,9 +1353,9 @@ def handle_realtime_stream_cost_calculation( potential_model_names = [] for result in results: if result["type"] == "session.created": - received_model = cast(OpenAIRealtimeStreamSessionEvents, result)["session"][ - "model" - ] + received_model = cast(OpenAIRealtimeStreamSessionEvents, result)[ + "session" + ].get("model", None) potential_model_names.append(received_model) potential_model_names.append(litellm_model_name) @@ -1364,6 +1364,8 @@ def handle_realtime_stream_cost_calculation( for model_name in potential_model_names: try: + if model_name is None: + continue _input_cost_per_token, _output_cost_per_token = generic_cost_per_token( model=model_name, usage=combined_usage_object, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 616d1a3db94..3b3e15cae16 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -5,7 +5,7 @@ from typing import Literal, Optional, Tuple, cast import litellm from litellm import verbose_logger -from litellm.types.utils import ModelInfo, Usage +from litellm.types.utils import CallTypes, ModelInfo, PassthroughCallTypes, Usage from litellm.utils import get_model_info @@ -343,3 +343,28 @@ def generic_cost_per_token( completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token return prompt_cost, completion_cost + + +class CostCalculatorUtils: + @staticmethod + def _call_type_has_image_response(call_type: str) -> bool: + """ + Returns True if the call type has an image response + + eg calls that have image response: + - Image Generation + - Image Edit + - Passthrough Image Generation + """ + if call_type in [ + # image generation + CallTypes.image_generation.value, + CallTypes.aimage_generation.value, + # passthrough image generation + PassthroughCallTypes.passthrough_image_generation.value, + # image edit + CallTypes.image_edit.value, + CallTypes.aimage_edit.value, + ]: + return True + return False diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 45951f24685..84b7a7997ee 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -219,6 +219,8 @@ class CallTypes(Enum): text_completion = "text_completion" image_generation = "image_generation" aimage_generation = "aimage_generation" + image_edit = "image_edit" + aimage_edit = "aimage_edit" moderation = "moderation" amoderation = "amoderation" atranscription = "atranscription" @@ -283,6 +285,8 @@ CallTypesLiteral = Literal[ "text_completion", "image_generation", "aimage_generation", + "image_edit", + "aimage_edit", "moderation", "amoderation", "atranscription", @@ -1600,6 +1604,8 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): Happens when their OpenAIImageResponse has the old OpenAI usage class. """ + model_config = ConfigDict(extra="allow", protected_namespaces=()) + def __init__( self, created: Optional[int] = None, @@ -1607,6 +1613,7 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): response_ms=None, usage: Optional[ImageUsage] = None, hidden_params: Optional[dict] = None, + **kwargs, ): if response_ms: _response_ms = response_ms diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 9e0ea5de07f..6c6df3627fa 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -2,6 +2,8 @@ import logging import os import sys import traceback +import asyncio +from typing import Optional import pytest import base64 from io import BytesIO @@ -14,6 +16,17 @@ sys.path.insert( import litellm from litellm.utils import ImageResponse +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardLoggingPayload + +class TestCustomLogger(CustomLogger): + def __init__(self): + self.standard_logging_payload: Optional[StandardLoggingPayload] = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_payload = kwargs.get("standard_logging_object", None) + pass + # Get the current directory of the file being run pwd = os.path.dirname(os.path.realpath(__file__)) @@ -241,3 +254,160 @@ async def test_azure_image_edit_litellm_sdk(): with open("test_image_edit.png", "wb") as f: f.write(image_bytes) + + +@pytest.mark.asyncio +async def test_openai_image_edit_cost_tracking(): + """Test OpenAI image edit cost tracking with custom logger""" + from litellm import image_edit, aimage_edit + test_custom_logger = TestCustomLogger() + litellm.logging_callback_manager._reset_all_callbacks() + litellm.callbacks = [test_custom_logger] + + # Mock response for Azure image edit + mock_response = { + "created": 1589478378, + "data": [ + { + "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" + } + ] + } + + class MockResponse: + def __init__(self, json_data, status_code): + self._json_data = json_data + self.status_code = status_code + self.text = json.dumps(json_data) + + def json(self): + return self._json_data + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + # Configure the mock to return our response + mock_post.return_value = MockResponse(mock_response, 200) + + litellm._turn_on_debug() + + prompt = """ + Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO. + """ + + # Set up test environment variables + + result = await aimage_edit( + prompt=prompt, + model="openai/gpt-image-1", + image=TEST_IMAGES, + ) + + # Verify the request was made correctly + mock_post.assert_called_once() + + + # Validate the response meets expected schema + ImageResponse.model_validate(result) + + if isinstance(result, ImageResponse) and result.data: + image_base64 = result.data[0].b64_json + if image_base64: + image_bytes = base64.b64decode(image_base64) + + # Save the image to a file + with open("test_image_edit.png", "wb") as f: + f.write(image_bytes) + + + await asyncio.sleep(5) + print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str)) + + # check model + assert test_custom_logger.standard_logging_payload["model"] == "gpt-image-1" + assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "openai" + + # check response_cost + assert test_custom_logger.standard_logging_payload["response_cost"] is not None + assert test_custom_logger.standard_logging_payload["response_cost"] > 0 + + + + +@pytest.mark.asyncio +async def test_azure_image_edit_cost_tracking(): + """Test Azure image edit cost tracking with custom logger""" + from litellm import image_edit, aimage_edit + test_custom_logger = TestCustomLogger() + litellm.logging_callback_manager._reset_all_callbacks() + litellm.callbacks = [test_custom_logger] + + # Mock response for Azure image edit + mock_response = { + "created": 1589478378, + "data": [ + { + "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" + } + ] + } + + class MockResponse: + def __init__(self, json_data, status_code): + self._json_data = json_data + self.status_code = status_code + self.text = json.dumps(json_data) + + def json(self): + return self._json_data + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + # Configure the mock to return our response + mock_post.return_value = MockResponse(mock_response, 200) + + litellm._turn_on_debug() + + prompt = """ + Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO. + """ + + # Set up test environment variables + + result = await aimage_edit( + prompt=prompt, + model="azure/CUSTOM_AZURE_DEPLOYMENT_NAME", + base_model="azure/gpt-image-1", + image=TEST_IMAGES, + ) + + # Verify the request was made correctly + mock_post.assert_called_once() + + + # Validate the response meets expected schema + ImageResponse.model_validate(result) + + if isinstance(result, ImageResponse) and result.data: + image_base64 = result.data[0].b64_json + if image_base64: + image_bytes = base64.b64decode(image_base64) + + # Save the image to a file + with open("test_image_edit.png", "wb") as f: + f.write(image_bytes) + + + await asyncio.sleep(5) + print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str)) + + # check model + assert test_custom_logger.standard_logging_payload["model"] == "CUSTOM_AZURE_DEPLOYMENT_NAME" + assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "azure" + + # check response_cost + assert test_custom_logger.standard_logging_payload["response_cost"] is not None + assert test_custom_logger.standard_logging_payload["response_cost"] > 0