fix(azure_ai): honour global drop_params for image params MAI cannot serve

get_optional_params_image_gen only forwarded the per-call drop_params flag
to provider configs, so litellm_settings drop_params: true never dropped
the n the MAI generations endpoint ignores. The MAI edits config also
advertised and forwarded size, which that endpoint ignores. Non-numeric
and non-positive n now surface as a 400 instead of a 500 or a pass-through.
This commit is contained in:
mateo-berri 2026-09-08 20:20:57 -07:00
parent 38a3de7641
commit 8e3052ff65
5 changed files with 83 additions and 96 deletions

View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final
import httpx
from httpx._types import RequestFiles
@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import (
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
@ -26,65 +25,8 @@ if TYPE_CHECKING:
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: Final[dict[str, Any]] = {}
supported_params: Final = 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: Final = {
"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}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string."
)
return ["prompt", "image", "model", "n"]
def validate_environment(
self,

View file

@ -151,7 +151,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
if k in supported_params:
if k == "size" and v:
self._map_size_param(v, optional_params, model)
elif k == "n" and v is not None and int(v) > self.MAX_IMAGES_PER_REQUEST:
elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST:
if not drop_params:
raise self._unsupported(
model,
@ -185,6 +185,14 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
def _unsupported(model: str, message: str) -> UnsupportedParamsError:
return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model)
def _image_count(self, n: object, model: str) -> int:
if isinstance(n, int):
return n
try:
return int(str(n))
except ValueError:
raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.")
def _map_size_param(self, size: str, optional_params: dict, model: str) -> None:
size_mapping: Final = {
"1024x1024": (1024, 1024),

View file

@ -3396,7 +3396,7 @@ def get_optional_params_image_gen(
non_default_params=non_default_params,
optional_params=optional_params,
model=model or "",
drop_params=drop_params if drop_params is not None else False,
drop_params=litellm.drop_params is True or drop_params is True,
)
elif (
custom_llm_provider == "openai"

View file

@ -6,6 +6,7 @@ import pytest
import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.azure_ai.image_edit import (
AzureFoundryMAIImageEditConfig,
get_azure_ai_image_edit_config,
@ -70,44 +71,48 @@ class TestAzureMAIImageEdit:
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},
def test_get_optional_params_image_edit_size_raises_400(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
with pytest.raises(litellm.UnsupportedParamsError, match="size") as exc_info:
ImageEditRequestUtils.get_optional_params_image_edit(
model="MAI-Image-2.5",
image_edit_provider_config=AzureFoundryMAIImageEditConfig(),
image_edit_optional_params={"size": "1024x1024", "n": 1},
)
assert exc_info.value.status_code == 400
def test_get_optional_params_image_edit_size_dropped_with_drop_params(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
optional_params = ImageEditRequestUtils.get_optional_params_image_edit(
model="MAI-Image-2.5",
image_edit_provider_config=AzureFoundryMAIImageEditConfig(),
image_edit_optional_params={"size": "1024x1024", "n": 1},
drop_params=True,
)
assert optional_params["size"] == "1792x1024"
assert "size" not in optional_params
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={},
def test_get_optional_params_image_edit_without_size_forwards_nothing_extra(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
optional_params = ImageEditRequestUtils.get_optional_params_image_edit(
model="MAI-Image-2.5",
drop_params=True,
image_edit_provider_config=AzureFoundryMAIImageEditConfig(),
image_edit_optional_params={},
)
assert optional_params["size"] == "1024x1024"
assert optional_params == {}
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_image_edit_size_surfaces_as_400(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
with pytest.raises(litellm.BadRequestError) as exc_info:
litellm.image_edit(
model="azure_ai/MAI-Image-2.5",
image=io.BytesIO(b"fake-image-bytes"),
prompt="Turn this into a studio product shot",
size="1024x1024",
api_key="test-key",
api_base="https://my-resource.services.ai.azure.com",
)
assert exc_info.value.status_code == 400
def test_transform_image_edit_request_uses_image_field(self):
config = AzureFoundryMAIImageEditConfig()
@ -117,14 +122,14 @@ class TestAzureMAIImageEdit:
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},
image_edit_optional_request_params={"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 "size" not in data
assert data["n"] == 1
assert len(files) == 1
assert files[0][0] == "image"

View file

@ -253,8 +253,8 @@ class TestAzureMAIImageGeneration:
assert optional_params["width"] == 1792
assert optional_params["height"] == 1024
@pytest.mark.parametrize("n", [2, 4, "2"])
def test_map_openai_params_multi_image_n_raises(self, n):
@pytest.mark.parametrize("n", [2, 4, "2", 0, -1])
def test_map_openai_params_n_other_than_one_raises(self, n):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"):
config.map_openai_params(
@ -264,6 +264,38 @@ class TestAzureMAIImageGeneration:
drop_params=False,
)
def test_map_openai_params_non_numeric_n_raises_400(self):
config = AzureFoundryMAIImageGenerationConfig()
with pytest.raises(UnsupportedParamsError, match="not a whole number of images") as exc_info:
config.map_openai_params(
non_default_params={"n": "abc"},
optional_params={},
model="MAI-Image-2.5",
drop_params=False,
)
assert exc_info.value.status_code == 400
def test_get_optional_params_image_gen_global_drop_params_drops_multi_image_n(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", True)
optional_params = get_optional_params_image_gen(
model="MAI-Image-2.5",
n=4,
custom_llm_provider="azure_ai",
provider_config=AzureFoundryMAIImageGenerationConfig(),
)
assert "n" not in optional_params
assert optional_params["width"] == 1024
def test_get_optional_params_image_gen_without_any_drop_params_still_raises(self, monkeypatch):
monkeypatch.setattr(litellm, "drop_params", False)
with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"):
get_optional_params_image_gen(
model="MAI-Image-2.5",
n=4,
custom_llm_provider="azure_ai",
provider_config=AzureFoundryMAIImageGenerationConfig(),
)
def test_map_openai_params_multi_image_n_dropped_with_drop_params(self):
config = AzureFoundryMAIImageGenerationConfig()
optional_params = config.map_openai_params(
@ -284,7 +316,7 @@ class TestAzureMAIImageGeneration:
)
assert optional_params["n"] == 1
@pytest.mark.parametrize("params", [{"n": 2}, {"size": "512x512"}, {"size": "1792x1024"}])
@pytest.mark.parametrize("params", [{"n": 2}, {"n": "abc"}, {"size": "512x512"}, {"size": "1792x1024"}])
def test_image_generation_rejected_params_surface_as_400(self, params):
with pytest.raises(litellm.BadRequestError) as exc_info:
litellm.image_generation(