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 = {