From 6d8fd2b5ee41b152dc68845bf0b311576a7fbc18 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 10:50:23 +0530 Subject: [PATCH] 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 --- .../image_generation/cost_calculator.py | 26 +++- .../image_generation/mai_transformation.py | 10 +- ...odel_prices_and_context_window_backup.json | 12 ++ model_prices_and_context_window.json | 5 + .../test_mai_image_generation.py | 123 ++++++++++++++++++ 5 files changed, 167 insertions(+), 9 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index b67de9cb70d..f8c876bb5be 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/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)}" + ) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index e17b1d687a9..f06aacd11c7 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ed6de4fa6b7..666d502a9e3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d5e35268519..666d502a9e3 100644 --- a/model_prices_and_context_window.json +++ b/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" ] 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 4cd9a8a78ad..d58f783cb3f 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,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