From 99a39c919f26ab90f53d85207dc9f937c16616d3 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 8 Jun 2026 20:35:07 +0000
Subject: [PATCH] fix(azure_ai): preserve zero num_output_tokens in MAI usage
 normalization

---
 .../azure_ai/image_generation/mai_transformation.py | 10 +++++-----
 .../image_generation/test_mai_image_generation.py   | 13 +++++++++++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py
index 04f7a0081e4..fd271ea3a17 100644
--- a/litellm/llms/azure_ai/image_generation/mai_transformation.py
+++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py
@@ -91,11 +91,11 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
 
         output_tokens = normalized_usage.get("output_tokens")
         if output_tokens is None:
-            output_tokens = (
-                normalized_usage.get("num_output_tokens")
-                or normalized_usage.get("output_image_tokens")
-                or 0
-            )
+            output_tokens = normalized_usage.get("num_output_tokens")
+        if output_tokens is None:
+            output_tokens = normalized_usage.get("output_image_tokens")
+        if output_tokens is None:
+            output_tokens = 0
 
         input_tokens = normalized_usage.get("input_tokens")
         if input_tokens is None:
diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
index c74c863b313..f7ca99f9716 100644
--- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
+++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
@@ -246,6 +246,19 @@ class TestAzureMAIImageGeneration:
         assert image_response.usage.input_tokens == 22
         assert image_response.usage.total_tokens == 1046
 
+    def test_normalize_mai_usage_preserves_zero_output_tokens(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        normalized = config.normalize_mai_image_usage(
+            {
+                "num_output_tokens": 0,
+                "output_image_tokens": 1024,
+                "num_input_text_tokens": 22,
+            }
+        )
+        assert normalized["output_tokens"] == 0
+        assert normalized["input_tokens"] == 22
+        assert normalized["total_tokens"] == 22
+
     def test_azure_sync_image_generation_uses_mai_response_transform(self):
         raw_response = MagicMock(spec=httpx.Response)
         raw_response.json.return_value = {