From f1082013fca18af515326adf3f0a7193feedd01b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:23:45 +0000 Subject: [PATCH] fix(azure_ai): build MAI image generation URL correctly for /mai/ root bases --- .../image_generation/mai_transformation.py | 16 ++++++++++------ .../test_mai_image_generation.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 1a004d885e6..071ca9d9895 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -28,15 +28,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if api_base is None: raise ValueError("api_base is required for Azure AI MAI image generation") - api_base = api_base.rstrip("/") api_version = api_version or "preview" + path, separator, query = api_base.partition("?") + path = path.rstrip("/") - if "/mai/" in api_base: - if "?" in api_base: - return api_base - return f"{api_base}?api-version={api_version}" + if "/mai/" in path: + prefix, _, _ = path.partition("/images/") + path = f"{prefix}/images/generations" + else: + path = f"{path}/mai/v1/images/generations" - return f"{api_base}/mai/v1/images/generations?api-version={api_version}" + if separator: + return f"{path}?{query}" + return f"{path}?api-version={api_version}" @staticmethod def get_mai_image_edit_url( 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 7afdf1f8499..f7ad333293c 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 @@ -80,6 +80,16 @@ class TestAzureMAIImageGeneration: ) 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)