From e4172c2e93e02d221bf788cf7e0f27301d741209 Mon Sep 17 00:00:00 2001
From: Emerson Gomes <emerson.gomes@thalesgroup.com>
Date: Sat, 2 May 2026 15:32:25 -0500
Subject: [PATCH] Break Azure AI image cost import cycle

---
 .../image_generation/cost_calculator.py       | 37 ++++++++++++++++---
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py
index 322d239751d..7fe05138899 100644
--- a/litellm/llms/azure_ai/image_generation/cost_calculator.py
+++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py
@@ -1,9 +1,6 @@
 from typing import Any
 
 import litellm
-from litellm.litellm_core_utils.llm_cost_calc.utils import (
-    calculate_image_response_cost_from_usage,
-)
 from litellm.types.utils import ImageResponse
 
 
@@ -23,10 +20,9 @@ def cost_calculator(
             f"image_response must be of type ImageResponse got type={type(image_response)}"
         )
 
-    usage_based_cost = calculate_image_response_cost_from_usage(
+    usage_based_cost = _calculate_cost_from_usage(
         model=model,
         image_response=image_response,
-        custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
     )
     if usage_based_cost is not None:
         return usage_based_cost
@@ -38,3 +34,34 @@ def cost_calculator(
     output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
     num_images = len(image_response.data) if image_response.data else 0
     return output_cost_per_image * num_images
+
+
+def _calculate_cost_from_usage(
+    model: str,
+    image_response: ImageResponse,
+) -> float | None:
+    usage = image_response.usage
+    if usage is None:
+        return None
+
+    prompt_tokens = usage.input_tokens
+    completion_tokens = usage.output_tokens
+    total_tokens = usage.total_tokens
+    if prompt_tokens is None or completion_tokens is None or total_tokens is None:
+        return None
+    if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
+        return None
+
+    model_info = litellm.get_model_info(
+        model=model,
+        custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
+    )
+    input_cost_per_token = model_info.get("input_cost_per_token") or 0.0
+    output_cost_per_image_token = (
+        model_info.get("output_cost_per_image_token")
+        or model_info.get("output_cost_per_token")
+        or 0.0
+    )
+    return (float(prompt_tokens) * input_cost_per_token) + (
+        float(completion_tokens) * output_cost_per_image_token
+    )