mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
[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
This commit is contained in:
Ishaan Jaff
GitHub
parent
95de41d10c
commit
a76d991952
4 changed files with 221 additions and 17 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue