From c1316a4d82ba3f026cb047ec4f443aa62a02a04b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 11:59:00 +0530 Subject: [PATCH] fix(azure_ai): normalize MAI usage in generation response handler Apply normalize_mai_image_usage before building ImageResponse so token-based cost calculation works when Azure returns num_output_tokens fields. Co-authored-by: Cursor --- .../image_generation/mai_transformation.py | 3 ++ .../test_mai_image_generation.py | 32 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 54ad06a9d41..e80a524e843 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -200,6 +200,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): ) -> ImageResponse: response = raw_response.json() + if "usage" in response: + response["usage"] = self.normalize_mai_image_usage(response.get("usage")) + logging_obj.post_call( input=request_data.get("prompt", ""), api_key=api_key, 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 7db8d7d265f..0143864f945 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 @@ -1,6 +1,8 @@ import os import sys +from unittest.mock import MagicMock +import httpx import pytest sys.path.insert(0, os.path.abspath("../../../../../..")) @@ -209,6 +211,36 @@ class TestAzureMAIImageGeneration: drop_params=False, ) + def test_transform_image_generation_response_normalizes_mai_usage(self): + config = AzureFoundryMAIImageGenerationConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "num_input_text_tokens": 22, + "output_image_tokens": 1024, + }, + } + + logging_obj = MagicMock() + image_response = config.transform_image_generation_response( + model="MAI-Image-2.5", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=logging_obj, + request_data={"prompt": "A red fox"}, + optional_params={"width": 1024, "height": 1024}, + litellm_params={}, + encoding=None, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.input_tokens == 22 + assert image_response.usage.total_tokens == 1046 + def test_mai_image_cost_calculator_token_based(self): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="")