fix(cost): treat ImageUsage as a known usage object in completion_cost

completion_cost flattened ImageUsage into chat Usage before routing to the
image cost calculator, dropping input_tokens_details. Image edit calls then
billed image input tokens at input_cost_per_token instead of
input_cost_per_image_token (MAI edits under-billed at $5/M instead of $8/M).
This commit is contained in:
mateo-berri 2026-07-28 11:11:58 -07:00
parent 88c2225e3d
commit 1d9090c3ee
2 changed files with 38 additions and 0 deletions

View file

@ -97,6 +97,7 @@ from litellm.types.utils import (
LiteLLMRealtimeStreamLoggingObject,
LlmProviders,
LlmProvidersSet,
ImageUsage,
ModelInfo,
ServiceTier,
StandardBuiltInToolsParams,
@ -913,6 +914,7 @@ def _is_known_usage_objects(usage_obj):
return (
isinstance(usage_obj, litellm.Usage)
or isinstance(usage_obj, ResponseAPIUsage)
or isinstance(usage_obj, ImageUsage)
or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj)
)

View file

@ -3509,3 +3509,39 @@ def test_combine_usage_objects_sums_mirrored_cache_write_fields_once():
assert combined_pair.prompt_tokens_details is not None
assert combined_pair.prompt_tokens_details.cache_write_tokens == 100
assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100
def test_completion_cost_image_edit_preserves_image_usage_input_split(monkeypatch):
"""
completion_cost must not flatten ImageUsage into chat Usage before routing to the
image cost calculator: doing so drops input_tokens_details, so image-input tokens
get billed at input_cost_per_token instead of input_cost_per_image_token
(e.g. azure_ai MAI image edits under-billed at $5/M instead of $8/M).
"""
from litellm.types.utils import (
ImageResponse,
ImageUsage,
ImageUsageInputTokensDetails,
)
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
response = ImageResponse(created=1785270000, data=[{"b64_json": "aGk="}])
response.usage = ImageUsage(
input_tokens=1061,
input_tokens_details=ImageUsageInputTokensDetails(image_tokens=1024, text_tokens=37),
output_tokens=1024,
total_tokens=2085,
)
cost = completion_cost(
completion_response=response,
model="azure_ai/MAI-Image-2.5-Pro",
custom_llm_provider="azure_ai",
call_type="aimage_edit",
)
expected = 37 * 5e-06 + 1024 * 8e-06 + 1024 * 0.000106
assert cost == pytest.approx(expected)
assert isinstance(response.usage, ImageUsage)