diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py new file mode 100644 index 0000000000..f476d6a94e --- /dev/null +++ b/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//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//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) diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index d471f496af..f3ae2d32ea 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/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: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7de88e30ce..6dfc68f724 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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), ) diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index f1421d4d57..c8a1e8f0e1 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/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: diff --git a/litellm/utils.py b/litellm/utils.py index f4f3707ebf..d2e1437c4b 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/image_gen_tests/test_image_edit.png b/tests/image_gen_tests/test_image_edit.png index c63c01cc7c..0f2de3749d 100644 Binary files a/tests/image_gen_tests/test_image_edit.png and b/tests/image_gen_tests/test_image_edit.png differ diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 649006dea0..9e0ea5de07 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/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) +