fix(azure_ai): address MAI image generation review feedback

Validate unsupported size values, default width/height independently, add MAI-Image-2.5 pricing, and expand test coverage.

@greptileai

Co-authored-by: Cursor <cursoragent@cursor.com>

litellm/llms/azure_ai/image_generation/cost_calculator.py

@@ -1,6 +1,9 @@
 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
 
 
@@ -9,19 +12,28 @@ def cost_calculator(
     image_response: Any,
 ) -> float:
     """
-    Recraft image generation cost calculator
+    Azure AI image generation cost calculator
     """
     _model_info = litellm.get_model_info(
         model=model,
         custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
     )
-    output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
-    num_images: int = 0
+
     if isinstance(image_response, ImageResponse):
+        token_based_cost = calculate_image_response_cost_from_usage(
+            model=model,
+            image_response=image_response,
+            custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
+        )
+        if token_based_cost is not None:
+            return token_based_cost
+
+        output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
+        num_images: int = 0
         if image_response.data:
             num_images = len(image_response.data)
         return output_cost_per_image * num_images
-    else:
-        raise ValueError(
-            f"image_response must be of type ImageResponse got type={type(image_response)}"
-        )
+
+    raise ValueError(
+        f"image_response must be of type ImageResponse got type={type(image_response)}"
+    )

litellm/llms/azure_ai/image_generation/mai_transformation.py

@@ -74,8 +74,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
                     f"Set drop_params=True to drop unsupported parameters."
                 )
 
-        if "width" not in optional_params and "height" not in optional_params:
+        if "width" not in optional_params:
             optional_params["width"] = self.DEFAULT_WIDTH
+        if "height" not in optional_params:
             optional_params["height"] = self.DEFAULT_HEIGHT
 
         optional_params.pop("size", None)
@@ -103,6 +104,11 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
                 raise ValueError(
                     f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
                 )
+        else:
+            raise ValueError(
+                f"Unsupported size value: '{size}'. "
+                f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
+            )
 
     def transform_image_generation_response(
         self,
@@ -134,5 +140,5 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
 
         width = optional_params.get("width", self.DEFAULT_WIDTH)
         height = optional_params.get("height", self.DEFAULT_HEIGHT)
-        image_response.size = f"{width}x{height}"
+        image_response.size = f"{width}x{height}"  # type: ignore[assignment]
         return image_response

litellm/model_prices_and_context_window_backup.json

@@ -6853,6 +6853,18 @@
             "/v1/images/generations"
         ]
     },
+    "azure_ai/MAI-Image-2.5": {
+        "input_cost_per_image_token": 8e-06,
+        "input_cost_per_token": 5e-06,
+        "litellm_provider": "azure_ai",
+        "mode": "image_generation",
+        "output_cost_per_image": 0.05,
+        "output_cost_per_image_token": 4.7e-05,
+        "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ]
+    },
     "azure_ai/Llama-3.2-11B-Vision-Instruct": {
         "input_cost_per_token": 3.7e-07,
         "litellm_provider": "azure_ai",

model_prices_and_context_window.json

@@ -6854,8 +6854,13 @@
         ]
     },
     "azure_ai/MAI-Image-2.5": {
+        "input_cost_per_image_token": 8e-06,
+        "input_cost_per_token": 5e-06,
         "litellm_provider": "azure_ai",
         "mode": "image_generation",
+        "output_cost_per_image": 0.05,
+        "output_cost_per_image_token": 4.7e-05,
+        "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632",
         "supported_endpoints": [
             "/v1/images/generations"
         ]

tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py

@@ -1,8 +1,11 @@
 import os
 import sys
 
+import pytest
+
 sys.path.insert(0, os.path.abspath("../../../../../.."))
 
+import litellm
 from litellm.llms.azure.azure import AzureChatCompletion
 from litellm.llms.azure.image_generation.http_utils import (
     azure_deployment_image_generation_json_body,
@@ -11,6 +14,10 @@ from litellm.llms.azure_ai.image_generation import (
     AzureFoundryMAIImageGenerationConfig,
     get_azure_ai_image_generation_config,
 )
+from litellm.llms.azure_ai.image_generation.cost_calculator import (
+    cost_calculator as azure_ai_image_cost_calculator,
+)
+from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails
 from litellm.utils import get_optional_params_image_gen
 
 
@@ -112,3 +119,119 @@ class TestAzureMAIImageGeneration:
         }
         out = azure_deployment_image_generation_json_body(api, data)
         assert out == data
+
+    def test_map_openai_params_custom_size(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        optional_params = config.map_openai_params(
+            non_default_params={"size": "768x768"},
+            optional_params={},
+            model="MAI-Image-2.5",
+            drop_params=True,
+        )
+        assert optional_params["width"] == 768
+        assert optional_params["height"] == 768
+
+    def test_map_openai_params_width_only_gets_height_default(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        optional_params = config.map_openai_params(
+            non_default_params={"width": 1792},
+            optional_params={},
+            model="MAI-Image-2.5",
+            drop_params=True,
+        )
+        assert optional_params["width"] == 1792
+        assert optional_params["height"] == config.DEFAULT_HEIGHT
+
+    def test_map_openai_params_height_only_gets_width_default(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        optional_params = config.map_openai_params(
+            non_default_params={"height": 1792},
+            optional_params={},
+            model="MAI-Image-2.5",
+            drop_params=True,
+        )
+        assert optional_params["width"] == config.DEFAULT_WIDTH
+        assert optional_params["height"] == 1792
+
+    def test_map_openai_params_unsupported_size_raises(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        with pytest.raises(ValueError, match="Unsupported size value: 'auto'"):
+            config.map_openai_params(
+                non_default_params={"size": "auto"},
+                optional_params={},
+                model="MAI-Image-2.5",
+                drop_params=True,
+            )
+
+    def test_map_openai_params_invalid_custom_size_raises(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"):
+            config.map_openai_params(
+                non_default_params={"size": "1024xabc"},
+                optional_params={},
+                model="MAI-Image-2.5",
+                drop_params=True,
+            )
+
+    def test_map_openai_params_unsupported_param_raises(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        with pytest.raises(ValueError, match="Parameter quality is not supported"):
+            config.map_openai_params(
+                non_default_params={"quality": "hd"},
+                optional_params={},
+                model="MAI-Image-2.5",
+                drop_params=False,
+            )
+
+    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="")
+        model = "azure_ai/MAI-Image-2.5"
+        model_info = litellm.get_model_info(
+            model=model, custom_llm_provider="azure_ai"
+        )
+        input_text_tokens = 100
+        output_image_tokens = 1024
+
+        image_response = ImageResponse(
+            data=[ImageObject(b64_json="img1")],
+            usage=ImageUsage(
+                input_tokens=input_text_tokens,
+                input_tokens_details=ImageUsageInputTokensDetails(
+                    text_tokens=input_text_tokens,
+                    image_tokens=0,
+                ),
+                output_tokens=output_image_tokens,
+                total_tokens=input_text_tokens + output_image_tokens,
+            ),
+        )
+
+        cost = azure_ai_image_cost_calculator(
+            model=model,
+            image_response=image_response,
+        )
+
+        expected_cost = (
+            input_text_tokens * model_info["input_cost_per_token"]
+            + output_image_tokens * model_info["output_cost_per_image_token"]
+        )
+        assert round(cost, 10) == round(expected_cost, 10)
+
+    def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self):
+        os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+        litellm.model_cost = litellm.get_model_cost_map(url="")
+        model = "azure_ai/MAI-Image-2.5"
+        model_info = litellm.get_model_info(
+            model=model, custom_llm_provider="azure_ai"
+        )
+        image_response = ImageResponse(
+            data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
+        )
+
+        cost = azure_ai_image_cost_calculator(
+            model=model,
+            image_response=image_response,
+        )
+
+        assert cost == len(image_response.data or []) * model_info["output_cost_per_image"]
+        assert cost > 0