ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(azure_ai): add MAI-Image-2.5 image generation support (#29688)

Browse Source 
* feat(azure_ai): add MAI-Image-2.5 image generation support

Route azure_ai MAI models to /mai/v1/images/generations and map OpenAI size to width/height for the serverless API.

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

* 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>

* feat(azure_ai): add MAI image edit and expand model cost map

Add MAI image edit support with usage normalization for Azure response format,
and register MAI-Image-2.5-Flash and MAI-Image-2e pricing in the model map.

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

* fix(azure_ai): validate MAI edit size by consuming map iterator

Greptile: lazy map() never evaluated int() so values like 1024xabc passed through.
Co-authored-by: Cursor <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

* fix(azure_ai): narrow MAI edit size param type for mypy

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

* Fix Azure MAI image response handling

* Fix MAI image generation base model routing

* fix(azure_ai): preserve zero num_output_tokens in MAI usage normalization

* fix(azure_ai): wrap MAI generation response JSON parsing in error handling

* fix(azure_ai): build MAI image edit URL correctly for /mai/ root bases

* fix(azure_ai): build MAI image generation URL correctly for /mai/ root bases

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Sameer Kankute 2026-06-09 06:57:04 +05:30 committed by GitHub
parent 92817cb65b
commit 424db6a980
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 1149 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
litellm/llms/azure/azure.py
View File

@ -43,7 +43,10 @@ from .common_utils import (
process_azure_headers,
select_azure_base_url_or_endpoint,
)
from .image_generation import get_azure_image_generation_config
from .image_generation import (
AzureFoundryMAIImageGenerationConfig,
get_azure_image_generation_config,
)
from .image_generation.http_utils import azure_deployment_image_generation_json_body
@ -1097,10 +1100,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
)
def create_azure_base_url(
self, azure_client_params: dict, model: Optional[str]
self,
azure_client_params: dict,
model: Optional[str],
base_model: Optional[str] = None,
) -> str:
from litellm.llms.azure_ai.image_generation import (
AzureFoundryFluxImageGenerationConfig,
AzureFoundryMAIImageGenerationConfig,
)
api_base: str = azure_client_params.get(
@ -1112,6 +1119,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if model is None:
model = ""
if AzureFoundryMAIImageGenerationConfig.is_mai_model(base_model or model):
return AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
api_base=api_base,
api_version=api_version,
)
# Handle FLUX 2 models on Azure AI which use a different URL pattern
# e.g., /providers/blackforestlabs/v1/flux-2-pro instead of /openai/deployments/{model}/images/generations
if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model):
@ -1153,10 +1166,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if api_base.endswith("/"):
api_base = api_base.rstrip("/")
api_version: str = azure_client_params.get("api_version", "")
# Use the deployment name (model) for URL construction, not the base_model from data
img_gen_api_base = self.create_azure_base_url(
azure_client_params=azure_client_params,
model=model or data.get("model", ""),
base_model=data.get("model", ""),
)
## LOGGING
@ -1285,9 +1298,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if aimg_generation is True:
return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers, model=model) # type: ignore
# Use the deployment name (model) for URL construction, not the base_model from data
img_gen_api_base = self.create_azure_base_url(
azure_client_params=azure_client_params, model=model
azure_client_params=azure_client_params,
model=model,
base_model=base_model,
)
## LOGGING
@ -1309,6 +1323,21 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
data=data,
headers=headers,
)
provider_config = get_azure_image_generation_config(
data.get("model", "dall-e-2")
)
if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig):
return provider_config.transform_image_generation_response(
model=data.get("model", "dall-e-2"),
raw_response=httpx_response,
model_response=model_response or ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=data,
litellm_params=data,
encoding=litellm.encoding,
)
response = httpx_response.json()
## LOGGING

3
litellm/llms/azure/image_generation/__init__.py
View File

@ -1,4 +1,5 @@
from litellm._logging import verbose_logger
from litellm.llms.azure_ai.image_generation import AzureFoundryMAIImageGenerationConfig
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
@ -24,6 +25,8 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig:
return AzureDallE2ImageGenerationConfig()
elif "dalle3" in model:
return AzureDallE3ImageGenerationConfig()
elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
return AzureFoundryMAIImageGenerationConfig()
else:
verbose_logger.debug(
f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image model format."

14
litellm/llms/azure_ai/image_edit/__init__.py
View File

@ -1,21 +1,33 @@
from litellm.llms.azure_ai.image_generation.flux_transformation import (
AzureFoundryFluxImageGenerationConfig,
)
from litellm.llms.azure_ai.image_generation.mai_transformation import (
AzureFoundryMAIImageGenerationConfig,
)
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from .flux2_transformation import AzureFoundryFlux2ImageEditConfig
from .mai_transformation import AzureFoundryMAIImageEditConfig
from .transformation import AzureFoundryFluxImageEditConfig
__all__ = ["AzureFoundryFluxImageEditConfig", "AzureFoundryFlux2ImageEditConfig"]
__all__ = [
"AzureFoundryFluxImageEditConfig",
"AzureFoundryFlux2ImageEditConfig",
"AzureFoundryMAIImageEditConfig",
]
def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig:
"""
Get the appropriate image edit config for an Azure AI model.
- MAI models use /mai/v1/images/edits with multipart form data and size
- FLUX 2 models use JSON with base64 image
- FLUX 1 models use multipart/form-data
"""
if AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
return AzureFoundryMAIImageEditConfig()
# Check if it's a FLUX 2 model
if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model):
return AzureFoundryFlux2ImageEditConfig()

199
litellm/llms/azure_ai/image_edit/mai_transformation.py Normal file
View File

@ -0,0 +1,199 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
import httpx
from httpx._types import RequestFiles
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
from litellm.llms.azure_ai.image_generation.mai_transformation import (
AzureFoundryMAIImageGenerationConfig,
)
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.llms.openai import FileTypes
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
from litellm.utils import convert_to_model_response_object
if TYPE_CHECKING:
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig):
"""Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5)."""
DEFAULT_SIZE = "1024x1024"
def get_supported_openai_params(self, model: str) -> list:
return ["prompt", "image", "model", "n", "size"]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> Dict:
optional_params: Dict[str, Any] = {}
supported_params = self.get_supported_openai_params(model)
for key, value in dict(image_edit_optional_params).items():
if value is None or key in optional_params:
continue
if key in supported_params:
if key == "size" and value:
size_param = cast(str, value)
self._validate_size_param(size_param)
optional_params[key] = size_param
else:
optional_params[key] = value
elif not drop_params:
raise ValueError(
f"Parameter {key} is not supported for model {model}. "
f"Supported parameters are {supported_params}. "
f"Set drop_params=True to drop unsupported parameters."
)
if "size" not in optional_params:
optional_params["size"] = self.DEFAULT_SIZE
return optional_params
def _validate_size_param(self, size: str) -> None:
known_sizes = {
"1024x1024",
"1792x1024",
"1024x1792",
"512x512",
"256x256",
}
if size in known_sizes:
return
if "x" in size:
try:
tuple(map(int, size.lower().split("x", 1)))
return
except ValueError:
raise ValueError(
f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
)
raise ValueError(
f"Unsupported size value: '{size}'. "
f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
)
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
api_key = AzureFoundryModelInfo.get_api_key(api_key)
if not api_key:
raise ValueError(
f"Azure AI API key is required for model {model}. "
"Set AZURE_AI_API_KEY environment variable or pass api_key parameter."
)
headers.update({"api-key": api_key})
return headers
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
api_base = AzureFoundryModelInfo.get_api_base(api_base)
if api_base is None:
raise ValueError(
"Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter."
)
api_version = (
litellm_params.get("api_version")
or get_secret_str("AZURE_AI_API_VERSION")
or "preview"
)
return AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url(
api_base=api_base,
api_version=api_version,
)
def transform_image_edit_request(
self,
model: str,
prompt: Optional[str],
image: Optional[FileTypes],
image_edit_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Tuple[Dict, RequestFiles]:
request_params = {
"model": model,
**image_edit_optional_request_params,
}
if prompt is not None:
request_params["prompt"] = prompt
data_without_files = {
key: value
for key, value in request_params.items()
if key not in ["image", "mask"]
}
files_list: List[Tuple[str, Any]] = []
if image is not None:
image_list = [image] if not isinstance(image, list) else image
for _image in image_list:
if _image is not None:
self._add_image_to_files(
files_list=files_list,
image=_image,
field_name="image",
)
break
return data_without_files, files_list
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> ImageResponse:
try:
response = raw_response.json()
except Exception:
raise OpenAIError(
message=raw_response.text, status_code=raw_response.status_code
)
if "usage" in response:
response["usage"] = (
AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage(
response.get("usage")
)
)
logging_obj.post_call(
input="",
api_key="",
additional_args={"complete_input_dict": {}},
original_response=response,
)
return convert_to_model_response_object(
response_object=response,
model_response_object=ImageResponse(),
response_type="image_generation",
)

4
litellm/llms/azure_ai/image_generation/__init__.py
View File

@ -7,12 +7,14 @@ from .dall_e_2_transformation import AzureFoundryDallE2ImageGenerationConfig
from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig
from .flux_transformation import AzureFoundryFluxImageGenerationConfig
from .gpt_transformation import AzureFoundryGPTImageGenerationConfig
from .mai_transformation import AzureFoundryMAIImageGenerationConfig
__all__ = [
"AzureFoundryFluxImageGenerationConfig",
"AzureFoundryGPTImageGenerationConfig",
"AzureFoundryDallE2ImageGenerationConfig",
"AzureFoundryDallE3ImageGenerationConfig",
"AzureFoundryMAIImageGenerationConfig",
]
@ -24,6 +26,8 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi
return AzureFoundryDallE2ImageGenerationConfig()
elif "dalle3" in model:
return AzureFoundryDallE3ImageGenerationConfig()
elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
return AzureFoundryMAIImageGenerationConfig()
elif "flux" in model:
return AzureFoundryFluxImageGenerationConfig()
else:

26
litellm/llms/azure_ai/image_generation/cost_calculator.py
View File

@ -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)}"
)

236
litellm/llms/azure_ai/image_generation/mai_transformation.py Normal file
View File

@ -0,0 +1,236 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.openai.common_utils import OpenAIError
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageResponse
from litellm.utils import convert_to_model_response_object
if TYPE_CHECKING:
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
"""Azure AI Foundry MAI image generation (e.g. MAI-Image-2.5)."""
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
@staticmethod
def get_mai_image_generation_url(
api_base: Optional[str],
api_version: Optional[str],
) -> str:
if api_base is None:
raise ValueError("api_base is required for Azure AI MAI image generation")
api_version = api_version or "preview"
path, separator, query = api_base.partition("?")
path = path.rstrip("/")
if "/mai/" in path:
prefix, _, _ = path.partition("/images/")
path = f"{prefix}/images/generations"
else:
path = f"{path}/mai/v1/images/generations"
if separator:
return f"{path}?{query}"
return f"{path}?api-version={api_version}"
@staticmethod
def get_mai_image_edit_url(
api_base: Optional[str],
api_version: Optional[str],
) -> str:
if api_base is None:
raise ValueError("api_base is required for Azure AI MAI image editing")
api_version = api_version or "preview"
path, separator, query = api_base.partition("?")
path = path.rstrip("/")
if "/mai/" in path:
prefix, _, _ = path.partition("/images/")
path = f"{prefix}/images/edits"
else:
path = f"{path}/mai/v1/images/edits"
if separator:
return f"{path}?{query}"
return f"{path}?api-version={api_version}"
@staticmethod
def is_mai_model(model: str) -> bool:
model_normalized = model.lower().replace("-", "").replace("_", "")
return "maiimage" in model_normalized
@staticmethod
def normalize_mai_image_usage(usage: Optional[Dict[str, Any]]) -> Dict[str, Any]:
"""Map Azure MAI usage fields to OpenAI ImageUsage schema."""
if usage is None:
return {
"input_tokens": 0,
"input_tokens_details": {"image_tokens": 0, "text_tokens": 0},
"output_tokens": 0,
"total_tokens": 0,
}
normalized_usage = dict(usage)
input_tokens_details = normalized_usage.get("input_tokens_details")
if not isinstance(input_tokens_details, dict):
input_tokens_details = {}
text_tokens = normalized_usage.get("num_input_text_tokens")
if text_tokens is None:
text_tokens = input_tokens_details.get("text_tokens")
if text_tokens is None:
text_tokens = normalized_usage.get("input_tokens", 0) or 0
image_tokens = normalized_usage.get("num_input_image_tokens")
if image_tokens is None:
image_tokens = input_tokens_details.get("image_tokens")
if image_tokens is None:
image_tokens = 0
output_tokens = normalized_usage.get("output_tokens")
if output_tokens is None:
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:
input_tokens = text_tokens + image_tokens
total_tokens = normalized_usage.get("total_tokens")
if total_tokens is None:
total_tokens = input_tokens + output_tokens
normalized_usage.update(
{
"input_tokens": input_tokens,
"input_tokens_details": {
"image_tokens": image_tokens,
"text_tokens": text_tokens,
},
"output_tokens": output_tokens,
"total_tokens": total_tokens,
}
)
return normalized_usage
def get_supported_openai_params(
self, model: str
) -> List[OpenAIImageGenerationOptionalParams]:
return ["n", "size"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params = self.get_supported_openai_params(model)
for k, v in non_default_params.items():
if k in optional_params:
continue
if k in supported_params:
if k == "size" and v:
self._map_size_param(v, optional_params)
else:
optional_params[k] = v
elif k in ("width", "height"):
optional_params[k] = v
elif not drop_params:
raise ValueError(
f"Parameter {k} is not supported for model {model}. "
f"Supported parameters are {supported_params} and width/height. "
f"Set drop_params=True to drop unsupported parameters."
)
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)
return optional_params
def _map_size_param(self, size: str, optional_params: dict) -> None:
size_mapping = {
"1024x1024": (1024, 1024),
"1792x1024": (1792, 1024),
"1024x1792": (1024, 1792),
"512x512": (512, 512),
"256x256": (256, 256),
}
if size in size_mapping:
width, height = size_mapping[size]
optional_params["width"] = width
optional_params["height"] = height
elif "x" in size:
try:
width, height = map(int, size.lower().split("x"))
optional_params["width"] = width
optional_params["height"] = height
except ValueError:
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,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
try:
response = raw_response.json()
except Exception:
raise OpenAIError(
message=raw_response.text, status_code=raw_response.status_code
)
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,
additional_args={"complete_input_dict": request_data},
original_response=response,
)
image_response: ImageResponse = convert_to_model_response_object(
response_object=response,
model_response_object=model_response,
response_type="image_generation",
)
width = optional_params.get("width", self.DEFAULT_WIDTH)
height = optional_params.get("height", self.DEFAULT_HEIGHT)
image_response.size = f"{width}x{height}" # type: ignore[assignment]
return image_response

37
litellm/model_prices_and_context_window_backup.json
View File

@ -6889,6 +6889,43 @@
"/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",
"/v1/images/edits"
]
},
"azure_ai/MAI-Image-2.5-Flash": {
"input_cost_per_image_token": 1.75e-06,
"input_cost_per_token": 1.75e-06,
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.0338,
"output_cost_per_image_token": 3.3e-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",
"/v1/images/edits"
]
},
"azure_ai/MAI-Image-2e": {
"input_cost_per_token": 5e-06,
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.02,
"output_cost_per_image_token": 1.95e-05,
"source": "https://aka.ms/mai-image-2e-foundryblog",
"supported_endpoints": [
"/v1/images/generations"
]
},
"azure_ai/Llama-3.2-11B-Vision-Instruct": {
"input_cost_per_token": 3.7e-07,
"litellm_provider": "azure_ai",

37
model_prices_and_context_window.json
View File

@ -6889,6 +6889,43 @@
"/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",
"/v1/images/edits"
]
},
"azure_ai/MAI-Image-2.5-Flash": {
"input_cost_per_image_token": 1.75e-06,
"input_cost_per_token": 1.75e-06,
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.0338,
"output_cost_per_image_token": 3.3e-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",
"/v1/images/edits"
]
},
"azure_ai/MAI-Image-2e": {
"input_cost_per_token": 5e-06,
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.02,
"output_cost_per_image_token": 1.95e-05,
"source": "https://aka.ms/mai-image-2e-foundryblog",
"supported_endpoints": [
"/v1/images/generations"
]
},
"azure_ai/Llama-3.2-11B-Vision-Instruct": {
"input_cost_per_token": 3.7e-07,
"litellm_provider": "azure_ai",

16
tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
View File

@ -57,6 +57,22 @@ def test_azure_providers_image_generation_json_body_keeps_model():
assert out == data
def test_azure_image_generation_mai_base_model_uses_mai_url():
azure_chat = AzureChatCompletion()
url = azure_chat.create_azure_base_url(
azure_client_params={
"azure_endpoint": "https://my-resource.services.ai.azure.com",
"api_version": "preview",
},
model="image-deployment-alias",
base_model="MAI-Image-2.5",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
)
def test_azure_image_generation_flattens_extra_body():
"""
Test that Azure image generation correctly flattens extra_body parameters.

171
tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py Normal file
View File

@ -0,0 +1,171 @@
import io
import os
import sys
from unittest.mock import MagicMock
import httpx
import pytest
sys.path.insert(0, os.path.abspath("../../../../../.."))
from litellm.llms.azure_ai.image_edit import (
AzureFoundryMAIImageEditConfig,
get_azure_ai_image_edit_config,
)
from litellm.llms.azure_ai.image_generation.mai_transformation import (
AzureFoundryMAIImageGenerationConfig,
)
class TestAzureMAIImageEdit:
def test_get_mai_image_edit_url(self):
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url(
api_base="https://my-resource.services.ai.azure.com",
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview"
)
def test_get_mai_image_edit_url_rewrites_generation_url(self):
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url(
api_base=(
"https://my-resource.services.ai.azure.com/mai/v1/images/generations"
"?api-version=preview"
),
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview"
)
def test_get_mai_image_edit_url_appends_edits_to_mai_root(self):
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url(
api_base="https://my-resource.services.ai.azure.com/mai/v1",
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview"
)
def test_get_azure_ai_image_edit_config_returns_mai(self):
config = get_azure_ai_image_edit_config("MAI-Image-2.5")
assert isinstance(config, AzureFoundryMAIImageEditConfig)
def test_validate_environment_uses_api_key_header(self):
config = AzureFoundryMAIImageEditConfig()
headers: dict = {}
config.validate_environment(headers, "MAI-Image-2.5", api_key="test-key")
assert headers["api-key"] == "test-key"
assert "Api-Key" not in headers
def test_get_complete_url(self):
config = AzureFoundryMAIImageEditConfig()
url = config.get_complete_url(
model="MAI-Image-2.5",
api_base="https://my-resource.services.ai.azure.com",
litellm_params={"api_version": "preview"},
)
assert "/mai/v1/images/edits" in url
assert "api-version=preview" in url
def test_map_openai_params_keeps_size(self):
config = AzureFoundryMAIImageEditConfig()
optional_params = config.map_openai_params(
image_edit_optional_params={"size": "1792x1024", "n": 1},
model="MAI-Image-2.5",
drop_params=True,
)
assert optional_params["size"] == "1792x1024"
assert optional_params["n"] == 1
assert "width" not in optional_params
assert "height" not in optional_params
def test_map_openai_params_defaults_size(self):
config = AzureFoundryMAIImageEditConfig()
optional_params = config.map_openai_params(
image_edit_optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert optional_params["size"] == "1024x1024"
def test_map_openai_params_unsupported_size_raises(self):
config = AzureFoundryMAIImageEditConfig()
with pytest.raises(ValueError, match="Unsupported size value: 'auto'"):
config.map_openai_params(
image_edit_optional_params={"size": "auto"},
model="MAI-Image-2.5",
drop_params=True,
)
def test_map_openai_params_invalid_size_format_raises(self):
config = AzureFoundryMAIImageEditConfig()
with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"):
config.map_openai_params(
image_edit_optional_params={"size": "1024xabc"},
model="MAI-Image-2.5",
drop_params=True,
)
def test_transform_image_edit_request_uses_image_field(self):
config = AzureFoundryMAIImageEditConfig()
image_bytes = io.BytesIO(b"fake-image-bytes")
data, files = config.transform_image_edit_request(
model="MAI-Image-2.5",
prompt="Turn this into a studio product shot",
image=image_bytes,
image_edit_optional_request_params={"size": "1024x1024", "n": 1},
litellm_params={},
headers={},
)
assert data["model"] == "MAI-Image-2.5"
assert data["prompt"] == "Turn this into a studio product shot"
assert data["size"] == "1024x1024"
assert data["n"] == 1
assert len(files) == 1
assert files[0][0] == "image"
assert files[0][0] != "image[]"
def test_normalize_mai_image_usage_maps_edit_response_fields(self):
usage = AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage(
{
"num_output_tokens": 1024,
"output_image_tokens": 1024,
}
)
assert usage["output_tokens"] == 1024
assert usage["input_tokens"] == 0
assert usage["total_tokens"] == 1024
assert usage["input_tokens_details"]["text_tokens"] == 0
assert usage["input_tokens_details"]["image_tokens"] == 0
def test_transform_image_edit_response_parses_mai_usage(self):
config = AzureFoundryMAIImageEditConfig()
raw_response = MagicMock(spec=httpx.Response)
raw_response.status_code = 200
raw_response.text = ""
raw_response.json.return_value = {
"created": 1780897477,
"data": [{"b64_json": "abc123"}],
"usage": {
"num_output_tokens": 1024,
"output_image_tokens": 1024,
},
}
logging_obj = MagicMock()
image_response = config.transform_image_edit_response(
model="MAI-Image-2.5",
raw_response=raw_response,
logging_obj=logging_obj,
)
assert image_response.data[0].b64_json == "abc123"
assert image_response.usage.output_tokens == 1024
assert image_response.usage.total_tokens == 1024

380
tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py Normal file
View File

@ -0,0 +1,380 @@
import os
import sys
from unittest.mock import MagicMock
import httpx
import pytest
sys.path.insert(0, os.path.abspath("../../../../../.."))
import litellm
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.image_generation import get_azure_image_generation_config
from litellm.llms.azure.image_generation.http_utils import (
azure_deployment_image_generation_json_body,
)
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
class TestAzureMAIImageGeneration:
def test_is_mai_model(self):
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5")
assert AzureFoundryMAIImageGenerationConfig.is_mai_model(
"azure_ai/MAI-Image-2.5"
)
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash")
assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e")
assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro")
assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1")
def test_mai_flash_and_2e_model_pricing_in_cost_map(self):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
flash_info = litellm.get_model_info(
model="azure_ai/MAI-Image-2.5-Flash",
custom_llm_provider="azure_ai",
)
assert flash_info["input_cost_per_token"] == 1.75e-06
assert flash_info["input_cost_per_image_token"] == 1.75e-06
assert flash_info["output_cost_per_image_token"] == 3.3e-05
image_2e_info = litellm.get_model_info(
model="azure_ai/MAI-Image-2e",
custom_llm_provider="azure_ai",
)
assert image_2e_info["input_cost_per_token"] == 5e-06
assert image_2e_info["output_cost_per_image_token"] == 1.95e-05
def test_get_mai_image_generation_url(self):
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
api_base="https://my-resource.services.ai.azure.com",
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
)
def test_get_mai_image_generation_url_preserves_full_path(self):
api = (
"https://my-resource.services.ai.azure.com/mai/v1/images/generations"
"?api-version=preview"
)
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
api_base=api,
api_version="preview",
)
assert url == api
def test_get_mai_image_generation_url_appends_generations_to_mai_root(self):
url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url(
api_base="https://my-resource.services.ai.azure.com/mai/v1",
api_version="preview",
)
assert (
url
== "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview"
)
def test_get_azure_ai_image_generation_config_returns_mai(self):
config = get_azure_ai_image_generation_config("MAI-Image-2.5")
assert isinstance(config, AzureFoundryMAIImageGenerationConfig)
def test_azure_image_generation_config_returns_mai(self):
config = get_azure_image_generation_config("MAI-Image-2.5")
assert isinstance(config, AzureFoundryMAIImageGenerationConfig)
def test_map_openai_params_size_to_width_height(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"size": "1024x1024", "n": 1},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert optional_params["width"] == 1024
assert optional_params["height"] == 1024
assert optional_params["n"] == 1
assert "size" not in optional_params
def test_map_openai_params_defaults(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert optional_params["width"] == 1024
assert optional_params["height"] == 1024
def test_get_optional_params_image_gen_mai(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = get_optional_params_image_gen(
model="MAI-Image-2.5",
size="1792x1024",
n=1,
custom_llm_provider="azure_ai",
provider_config=config,
drop_params=True,
)
assert optional_params["width"] == 1792
assert optional_params["height"] == 1024
assert "size" not in optional_params
def test_azure_create_azure_base_url_mai(self):
azure_chat = AzureChatCompletion()
url = azure_chat.create_azure_base_url(
azure_client_params={
"azure_endpoint": "https://my-resource.services.ai.azure.com",
"api_version": "preview",
},
model="MAI-Image-2.5",
)
assert "/mai/v1/images/generations" in url
assert "api-version=preview" in url
def test_mai_json_body_keeps_model(self):
api = (
"https://my-resource.services.ai.azure.com/mai/v1/images/generations"
"?api-version=preview"
)
data = {
"model": "MAI-Image-2.5",
"prompt": "A photograph of a red fox",
"width": 1024,
"height": 1024,
"n": 1,
}
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_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_transform_image_generation_response_non_json_raises_openai_error(self):
from litellm.llms.openai.common_utils import OpenAIError
config = AzureFoundryMAIImageGenerationConfig()
raw_response = MagicMock(spec=httpx.Response)
raw_response.json.side_effect = ValueError("not json")
raw_response.text = "upstream gateway error"
raw_response.status_code = 502
with pytest.raises(OpenAIError) as exc_info:
config.transform_image_generation_response(
model="MAI-Image-2.5",
raw_response=raw_response,
model_response=ImageResponse(),
logging_obj=MagicMock(),
request_data={"prompt": "A red fox"},
optional_params={"width": 1024, "height": 1024},
litellm_params={},
encoding=None,
)
assert exc_info.value.status_code == 502
assert exc_info.value.message == "upstream gateway error"
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 = {
"created": 1780897477,
"data": [{"b64_json": "abc123"}],
"usage": {
"num_output_tokens": 1024,
"num_input_text_tokens": 22,
},
}
class MAIImageGenerationAzureChatCompletion(AzureChatCompletion):
def make_sync_azure_httpx_request(self, **kwargs):
return raw_response
logging_obj = MagicMock()
image_response = MAIImageGenerationAzureChatCompletion().image_generation(
prompt="A red fox",
timeout=60.0,
optional_params={"width": 1792, "height": 1024},
logging_obj=logging_obj,
headers={},
model="MAI-Image-2.5",
api_key="test-key",
api_base="https://my-resource.services.ai.azure.com",
api_version="preview",
litellm_params={},
)
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
assert image_response.size == "1792x1024"
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