fix(azure_ai): stop accepting MAI image params the endpoint cannot honour

Two params were advertised for the MAI image models and dropped downstream,
so the caller got a 200 that did not match the request, or an opaque
provider 400.

n: get_supported_openai_params returns ["n", "size"], so n passes validation
and is forwarded. The MAI endpoint (/mai/v1/images/generations) has no count
field at all — its documented body is model/prompt/width/height, plus image
for edits — and ignores both `n` and the native `sampleCount`. Measured
against MAI-Image-2.5 and MAI-Image-2.5-Flash: n=2 and n=4 each return HTTP
200 with exactly one image, billed as one, with nothing in the response
saying the request was reduced. A caller balancing cost against image count
cannot see it. n=1 still passes through; n>1 now raises unless drop_params
is set, which is the existing opt-in for silently dropping a param.

size: _map_size_param's table offered five sizes, of which one is usable.
MAI requires width and height >= 768px and width*height <= 1048576, so
512x512 and 256x256 are under the per-side minimum and 1792x1024 / 1024x1792
are over the pixel budget — all four 400 at the provider with "Model does
not support request parameter value supplied: 'width' must be at least 768
pixels." Only 1024x1024 works. The bounds are now checked where the size is
mapped, so the error names the constraint instead of arriving from Azure.

width/height are deliberately left unchecked: they pass through unmapped, so
a future MAI model with different bounds stays reachable without a code
change.

Verified on a live Azure AI Foundry deployment of MAI-Image-2.5 and
MAI-Image-2.5-Flash (2026-08-17). One existing test asserted the 1792x1024
mapping; its size is changed to a size the provider accepts, keeping what it
was testing.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Mihidum Hettiyahandi 2026-08-21 09:42:08 +10:00
parent 168a0055a2
commit 808787659b
2 changed files with 122 additions and 33 deletions

View file

@ -21,6 +21,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
# The MAI endpoint produces exactly one image per request. Its documented
# body is model/prompt/width/height (plus `image` for edits) — there is no
# count field, and `n` (or the native `sampleCount`) is accepted and
# ignored, so a request for more silently comes back with one.
MAX_IMAGES_PER_REQUEST: Final = 1
# Provider-side bounds on the generated image. Both are enforced by the
# MAI endpoint, which 400s with "'width' must be at least 768 pixels."
# Only `size` is checked against them: `width`/`height` pass through
# unmapped, which keeps a future model with different bounds reachable.
MIN_DIMENSION_PX: Final = 768
MAX_TOTAL_PX: Final = 1024 * 1024
@staticmethod
def get_mai_image_generation_url(
api_base: str | None,
@ -146,6 +159,15 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
if k in supported_params:
if k == "size" and v:
self._map_size_param(v, optional_params)
elif k == "n" and v is not None and v > self.MAX_IMAGES_PER_REQUEST:
if not drop_params:
raise ValueError(
f"n={v} is not supported for model {model}. The Azure AI MAI image "
f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per "
"request and ignores any count, so a larger value would silently "
"return fewer images than requested. Send one request per image, or "
"set drop_params=True to drop n."
)
else:
optional_params[k] = v
elif k in ("width", "height"):
@ -176,13 +198,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
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:
@ -191,6 +209,29 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
)
self._validate_dimensions(size=size, width=width, height=height)
optional_params["width"] = width
optional_params["height"] = height
def _validate_dimensions(self, size: str, width: int, height: int) -> None:
"""Reject a `size` the MAI endpoint would 400 on.
Several OpenAI-standard sizes are outside MAI's bounds: 512x512 and
256x256 fall under the per-side minimum, and 1792x1024 / 1024x1792
exceed the total pixel budget. Checking here turns an opaque provider
400 into an error that names the constraint.
"""
if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX:
raise ValueError(
f"Unsupported size value: '{size}'. Azure AI MAI image models require width and "
f"height of at least {self.MIN_DIMENSION_PX} pixels."
)
if width * height > self.MAX_TOTAL_PX:
raise ValueError(
f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most "
f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height})."
)
def transform_image_generation_response(
self,
model: str,

View file

@ -1,10 +1,8 @@
import os
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.image_generation import get_azure_image_generation_config
@ -30,9 +28,7 @@ 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("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")
@ -62,16 +58,10 @@ class TestAzureMAIImageGeneration:
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"
)
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"
)
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",
@ -83,10 +73,7 @@ class TestAzureMAIImageGeneration:
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"
)
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")
@ -124,13 +111,13 @@ class TestAzureMAIImageGeneration:
config = AzureFoundryMAIImageGenerationConfig()
optional_params = get_optional_params_image_gen(
model="MAI-Image-2.5",
size="1792x1024",
size="1024x1024",
n=1,
custom_llm_provider="azure_ai",
provider_config=config,
drop_params=True,
)
assert optional_params["width"] == 1792
assert optional_params["width"] == 1024
assert optional_params["height"] == 1024
assert "size" not in optional_params
@ -147,10 +134,7 @@ class TestAzureMAIImageGeneration:
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"
)
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",
@ -214,6 +198,74 @@ class TestAzureMAIImageGeneration:
drop_params=True,
)
@pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"])
def test_map_openai_params_size_below_minimum_dimension_raises(self, size):
"""MAI requires >= 768px per side; the OpenAI size table offered smaller ones."""
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="at least 768 pixels"):
config.map_openai_params(
non_default_params={"size": size},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
@pytest.mark.parametrize("size", ["1792x1024", "1024x1792"])
def test_map_openai_params_size_over_total_pixel_budget_raises(self, size):
"""MAI caps total pixels at 1024*1024, so both landscape/portrait sizes 400 upstream."""
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="at most 1048576 total pixels"):
config.map_openai_params(
non_default_params={"size": size},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
def test_map_openai_params_explicit_width_height_not_range_checked(self):
"""width/height pass through unmapped, so a future model's bounds stay reachable."""
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"width": 1792, "height": 1024},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert optional_params["width"] == 1792
assert optional_params["height"] == 1024
@pytest.mark.parametrize("n", [2, 4])
def test_map_openai_params_multi_image_n_raises(self, n):
"""The MAI endpoint returns one image and ignores any count, so n>1 must not pass silently."""
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="returns exactly 1 image per request"):
config.map_openai_params(
non_default_params={"n": n},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
def test_map_openai_params_multi_image_n_dropped_with_drop_params(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"n": 4},
optional_params={},
model="MAI-Image-2.5",
drop_params=True,
)
assert "n" not in optional_params
def test_map_openai_params_single_image_n_still_passes_through(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
non_default_params={"n": 1},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
assert optional_params["n"] == 1
def test_map_openai_params_unsupported_param_raises(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(ValueError, match="Parameter quality is not supported"):
@ -363,16 +415,12 @@ class TestAzureMAIImageGeneration:
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")]
)
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 == len(image_response.data or []) * model_info["output_cost_per_image"]
assert cost > 0