Fix MAI image generation base model routing

cursoragent · cursoragent · commit d9e4eaab505c · 2026-06-08T07:42:19.000Z
diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py
@@ -1100,7 +1100,10 @@ def make_sync_azure_httpx_request(
         )
 
     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,
@@ -1116,8 +1119,7 @@ def create_azure_base_url(
         if model is None:
             model = ""
 
-        # MAI image models: /mai/v1/images/generations (serverless Azure AI)
-        if AzureFoundryMAIImageGenerationConfig.is_mai_model(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,
@@ -1164,10 +1166,10 @@ async def aimage_generation(
             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
@@ -1296,9 +1298,10 @@ def image_generation(
             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
diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
@@ -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.
