From 1d9090c3ee4a7a257173acbd38af38f77ae54fe8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 28 Jul 2026 11:11:58 -0700
Subject: [PATCH] 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).
---
 litellm/cost_calculator.py                 |  2 ++
 tests/test_litellm/test_cost_calculator.py | 36 ++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 96aed20529f..0a217c941b7 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -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)
     )
 
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 276ee96ed65..31b0cd67f25 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -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)