[Feat] Add /image/edits support for Azure (#11160)

* feat: add image edits on litellm router

* feat: add image edits endpoint

* fix: use pure async for image edits

* fix: base_process_llm_request

* fix: get_image_content_type

* feat: add image edits endpoint

* add image edits on UI

* test: image edits support

* fix: linting errors

* fix: linting errors

* test fix img gen

* feat: azure image edits

* fix: fix url construction of azure image edits

* fix: mock endpoints for azure images

litellm/llms/azure/image_edit/transformation.py

@@ -0,0 +1,83 @@
+from typing import Optional, cast
+
+import httpx
+
+import litellm
+from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.utils import _add_path_to_api_base
+
+
+class AzureImageEditConfig(OpenAIImageEditConfig):
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+    ) -> dict:
+        api_key = (
+            api_key
+            or litellm.api_key
+            or litellm.azure_key
+            or get_secret_str("AZURE_OPENAI_API_KEY")
+            or get_secret_str("AZURE_API_KEY")
+        )
+
+        headers.update(
+            {
+                "Authorization": f"Bearer {api_key}",
+            }
+        )
+        return headers
+
+    def get_complete_url(
+        self,
+        model: str,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        """
+        Constructs a complete URL for the API request.
+
+        Args:
+        - api_base: Base URL, e.g.,
+            "https://litellm8397336933.openai.azure.com"
+            OR
+            "https://litellm8397336933.openai.azure.com/openai/deployments/<deployment_name>/images/edits?api-version=2024-05-01-preview"
+        - model: Model name (deployment name).
+        - litellm_params: Additional query parameters, including "api_version".
+
+        Returns:
+        - A complete URL string, e.g.,
+        "https://litellm8397336933.openai.azure.com/openai/deployments/<deployment_name>/images/edits?api-version=2024-05-01-preview"
+        """
+        api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
+        if api_base is None:
+            raise ValueError(
+                f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`"
+            )
+        original_url = httpx.URL(api_base)
+
+        # Extract api_version or use default
+        api_version = cast(Optional[str], litellm_params.get("api_version"))
+
+        # Create a new dictionary with existing params
+        query_params = dict(original_url.params)
+
+        # Add api_version if needed
+        if "api-version" not in query_params and api_version:
+            query_params["api-version"] = api_version
+
+        # Add the path to the base URL using the model as deployment name
+        if "/openai/deployments/" not in api_base:
+            new_url = _add_path_to_api_base(
+                api_base=api_base,
+                ending_path=f"/openai/deployments/{model}/images/edits",
+            )
+        else:
+            new_url = api_base
+
+        # Use the new query_params dictionary
+        final_url = httpx.URL(new_url).copy_with(params=query_params)
+
+        return str(final_url)

litellm/llms/base_llm/image_edit/transformation.py

@@ -73,6 +73,7 @@ class BaseImageEditConfig(ABC):
     @abstractmethod
     def get_complete_url(
         self,
+        model: str,
         api_base: Optional[str],
         litellm_params: dict,
     ) -> str:

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2166,6 +2166,7 @@ class BaseLLMHTTPHandler:
             headers.update(extra_headers)
 
         api_base = image_edit_provider_config.get_complete_url(
+            model=model,
             api_base=litellm_params.api_base,
             litellm_params=dict(litellm_params),
         )
@@ -2250,6 +2251,7 @@ class BaseLLMHTTPHandler:
             headers.update(extra_headers)
 
         api_base = image_edit_provider_config.get_complete_url(
+            model=model,
             api_base=litellm_params.api_base,
             litellm_params=dict(litellm_params),
         )

litellm/llms/openai/image_edit/transformation.py

@@ -135,6 +135,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
 
     def get_complete_url(
         self,
+        model: str,
         api_base: Optional[str],
         litellm_params: dict,
     ) -> str:

litellm/utils.py

@@ -6684,6 +6684,12 @@ class ProviderConfigManager:
             )
 
             return OpenAIImageEditConfig()
+        if LlmProviders.AZURE == provider:
+            from litellm.llms.azure.image_edit.transformation import (
+                AzureImageEditConfig,
+            )
+
+            return AzureImageEditConfig()
         return None
 
 

tests/image_gen_tests/test_image_edits.py

@@ -5,6 +5,8 @@ import traceback
 import pytest
 import base64
 from io import BytesIO
+from unittest.mock import patch, AsyncMock
+import json
 
 sys.path.insert(
     0, os.path.abspath("../..")
@@ -143,3 +145,99 @@ async def test_openai_image_edit_with_bytesio():
                     f.write(image_bytes)
     except litellm.ContentPolicyViolationError as e:
         pass
+
+
+@pytest.mark.asyncio
+async def test_azure_image_edit_litellm_sdk():
+    """Test Azure image edit with mocked httpx request to validate request body and URL"""
+    from litellm import image_edit, aimage_edit
+    
+    # Mock response for Azure image edit
+    mock_response = {
+        "created": 1589478378,
+        "data": [
+            {
+                "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+            }
+        ]
+    }
+
+    class MockResponse:
+        def __init__(self, json_data, status_code):
+            self._json_data = json_data
+            self.status_code = status_code
+            self.text = json.dumps(json_data)
+
+        def json(self):
+            return self._json_data
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new_callable=AsyncMock,
+    ) as mock_post:
+        # Configure the mock to return our response
+        mock_post.return_value = MockResponse(mock_response, 200)
+
+        litellm._turn_on_debug()
+        
+        prompt = """
+        Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
+        """
+        
+        # Set up test environment variables
+        test_api_base = "https://ai-api-gw-uae-north.openai.azure.com"
+        test_api_key = "test-api-key"
+        test_api_version = "2025-04-01-preview"
+        
+        result = await aimage_edit(
+            prompt=prompt,
+            model="azure/gpt-image-1",
+            api_base=test_api_base,
+            api_key=test_api_key,
+            api_version=test_api_version,
+            image=TEST_IMAGES,
+        )
+        
+        # Verify the request was made correctly
+        mock_post.assert_called_once()
+        
+        # Check the URL
+        call_args = mock_post.call_args
+        expected_url = f"{test_api_base}/openai/deployments/gpt-image-1/images/edits?api-version={test_api_version}"
+        actual_url = call_args.args[0] if call_args.args else call_args.kwargs.get('url')
+        print(f"Expected URL: {expected_url}")
+        print(f"Actual URL: {actual_url}")
+        assert actual_url == expected_url, f"URL mismatch. Expected: {expected_url}, Got: {actual_url}"
+        
+        # Check the request body
+        if 'data' in call_args.kwargs:
+            # For multipart form data, check the data parameter
+            form_data = call_args.kwargs['data']
+            print("Form data keys:", list(form_data.keys()) if hasattr(form_data, 'keys') else "Not a dict")
+            
+            # Validate that model and prompt are in the form data
+            assert 'model' in form_data, "model should be in form data"
+            assert 'prompt' in form_data, "prompt should be in form data"
+            assert form_data['model'] == 'gpt-image-1', f"Expected model 'gpt-image-1', got {form_data['model']}"
+            assert prompt.strip() in form_data['prompt'], f"Expected prompt to contain '{prompt.strip()}'"
+            
+        # Check headers
+        headers = call_args.kwargs.get('headers', {})
+        print("Request headers:", headers)
+        assert 'Authorization' in headers, "Authorization header should be present"
+        assert headers['Authorization'].startswith('Bearer '), "Authorization should be Bearer token"
+        
+        print("result from image edit", result)
+
+        # Validate the response meets expected schema
+        ImageResponse.model_validate(result)
+        
+        if isinstance(result, ImageResponse) and result.data:
+            image_base64 = result.data[0].b64_json
+            if image_base64:
+                image_bytes = base64.b64decode(image_base64)
+
+                # Save the image to a file
+                with open("test_image_edit.png", "wb") as f:
+                    f.write(image_bytes)
+