From 92a917fc1cb8a452323a2b491f3b8d8688c59db8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:13:12 +0000 Subject: [PATCH] fix(azure_ai): build MAI image edit URL correctly for /mai/ root bases --- .../image_generation/mai_transformation.py | 17 ++++++++++------- .../test_mai_image_edit_transformation.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index c08d44f9031..1a004d885e6 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -46,16 +46,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if api_base is None: raise ValueError("api_base is required for Azure AI MAI image editing") - 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: - api_base = api_base.replace("/images/generations", "/images/edits") - 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/edits" + else: + path = f"{path}/mai/v1/images/edits" - return f"{api_base}/mai/v1/images/edits?api-version={api_version}" + if separator: + return f"{path}?{query}" + return f"{path}?api-version={api_version}" @staticmethod def is_mai_model(model: str) -> bool: diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 62126db2475..d5256be02d7 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -41,6 +41,16 @@ class TestAzureMAIImageEdit: == "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)