diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index 376bfa994de..ac2fe2cf605 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -1,13 +1,10 @@ -from typing import Dict, Optional, Tuple, cast +from typing import Dict, Optional, cast import httpx -from httpx._types import RequestFiles import litellm from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import FileTypes -from litellm.types.router import GenericLiteLLMParams from litellm.utils import _add_path_to_api_base @@ -100,32 +97,7 @@ class AzureImageEditConfig(OpenAIImageEditConfig): return str(final_url) - def transform_image_edit_request( - self, - model: str, - prompt: Optional[str], - image: Optional[FileTypes], - image_edit_optional_request_params: Dict, - litellm_params: GenericLiteLLMParams, - headers: dict, - ) -> Tuple[Dict, RequestFiles]: - data, files = super().transform_image_edit_request( - model=model, - prompt=prompt, - image=image, - image_edit_optional_request_params=image_edit_optional_request_params, - litellm_params=litellm_params, - headers=headers, - ) - litellm_params_dict = ( - litellm_params.model_dump(exclude_none=True) - if hasattr(litellm_params, "model_dump") - else dict(litellm_params) - ) - resolved_url = self.get_complete_url( - model=model, - api_base=litellm_params_dict.get("api_base"), - litellm_params=litellm_params_dict, - ) - data = self.azure_deployment_image_edit_form_data(data, resolved_url) - return data, files + def finalize_image_edit_multipart_data( + self, data: dict, resolved_request_url: str + ) -> dict: + return self.azure_deployment_image_edit_form_data(data, resolved_request_url) diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index cea96bde74d..71710f5b669 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -102,6 +102,15 @@ class BaseImageEditConfig(ABC): ) -> Tuple[Dict, RequestFiles]: pass + def finalize_image_edit_multipart_data( + self, data: dict, resolved_request_url: str + ) -> dict: + """ + Adjust non-file form fields after ``transform_image_edit_request`` using the + exact URL that will be used for the HTTP POST (same string as ``get_complete_url``). + """ + return data + @abstractmethod def transform_image_edit_response( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 0c4816fcda2..de476c2fb0d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5579,6 +5579,9 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + data = image_edit_provider_config.finalize_image_edit_multipart_data( + data, api_base + ) ## LOGGING logging_obj.pre_call( @@ -5677,6 +5680,9 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + data = image_edit_provider_config.finalize_image_edit_multipart_data( + data, api_base + ) ## LOGGING logging_obj.pre_call( diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index bfdd6b95f2f..d80e0aa90ad 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py @@ -20,7 +20,8 @@ def test_azure_deployment_image_edit_form_data_keeps_model_non_deployment_url(): assert out == data -def test_azure_transform_image_edit_request_omits_model_for_deployment(): +def test_azure_finalize_image_edit_strips_model_after_openai_transform(): + """OpenAI transform still includes model; finalize uses the real request URL.""" config = AzureImageEditConfig() model = "gpt-image-2-dep" prompt = "add a hat" @@ -37,7 +38,14 @@ def test_azure_transform_image_edit_request_omits_model_for_deployment(): litellm_params=litellm_params, headers={}, ) - assert "model" not in data - assert data.get("prompt") == prompt - assert data.get("n") == 1 + assert data.get("model") == model + resolved = config.get_complete_url( + model=model, + api_base=litellm_params.api_base, + litellm_params=litellm_params.model_dump(exclude_none=True), + ) + data_out = config.finalize_image_edit_multipart_data(data, resolved) + assert "model" not in data_out + assert data_out.get("prompt") == prompt + assert data_out.get("n") == 1 assert len(files) >= 1 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 index aec0dadd7d4..6e91fdf8653 100644 --- 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 @@ -12,6 +12,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.azure.image_generation import ( AzureDallE3ImageGenerationConfig, get_azure_image_generation_config, @@ -312,26 +313,21 @@ def test_azure_image_generation_base_model_vs_deployment_name(): optional_params = {"n": 1, "size": "1024x1024"} - # Mock the HTTP request to capture what gets sent + mock_http_response = MagicMock() + mock_http_response.status_code = 200 + mock_http_response.json.return_value = { + "created": 1234567890, + "data": [{"url": "https://example.com/image.png", "revised_prompt": prompt}], + } + with patch.object( - azure_chat_completion, - "make_sync_azure_httpx_request", - return_value=MagicMock( - json=lambda: { - "created": 1234567890, - "data": [ - {"url": "https://example.com/image.png", "revised_prompt": prompt} - ], - } - ), - ) as mock_request: - # Mock logging object + HTTPHandler, "post", return_value=mock_http_response + ) as mock_post: logging_obj = MagicMock() logging_obj.pre_call = MagicMock() logging_obj.post_call = MagicMock() - # Call the image_generation method - response = azure_chat_completion.image_generation( + azure_chat_completion.image_generation( prompt=prompt, timeout=60.0, optional_params=optional_params, @@ -344,34 +340,21 @@ def test_azure_image_generation_base_model_vs_deployment_name(): litellm_params=litellm_params, ) - # Verify the mock was called - assert mock_request.called, "HTTP request should have been made" - - # Get the call arguments - call_kwargs = mock_request.call_args.kwargs - - # Verify the URL uses the deployment name (not base_model) - api_base_used = call_kwargs.get("api_base", "") - assert model in api_base_used, ( - f"URL should contain deployment name '{model}', " - f"but got: {api_base_used}" - ) - assert base_model not in api_base_used or base_model == model, ( + assert mock_post.called, "HTTPHandler.post should be invoked" + post_kwargs = mock_post.call_args.kwargs + url_used = post_kwargs.get("url", "") + assert ( + model in url_used + ), f"URL should contain deployment name '{model}', but got: {url_used}" + assert base_model not in url_used or base_model == model, ( f"URL should NOT contain base_model '{base_model}' when it differs from deployment name, " - f"but got: {api_base_used}" + f"but got: {url_used}" ) - # Verify the HTTP JSON body omits model (deployment is only in the URL) - request_data = call_kwargs.get("data", {}) - wire_json = AzureChatCompletion.azure_deployment_image_generation_json_body( - api_base_used, request_data - ) + wire_json = post_kwargs.get("json") or {} assert ( "model" not in wire_json ), f"Azure deployment image gen must not send 'model' in JSON body; got keys: {list(wire_json)}" - assert request_data.get("model") == base_model # internal dict unchanged - - # Verify other fields are correct on the wire payload assert wire_json.get("prompt") == prompt assert wire_json.get("n") == 1 assert wire_json.get("size") == "1024x1024" @@ -402,27 +385,24 @@ async def test_azure_aimage_generation_base_model_vs_deployment_name(): "api_version": api_version, } - # Mock the HTTP request to capture what gets sent - with patch.object( - azure_chat_completion, - "make_async_azure_httpx_request", - new_callable=AsyncMock, - return_value=MagicMock( - json=lambda: { - "created": 1234567890, - "data": [ - {"url": "https://example.com/image.png", "revised_prompt": prompt} - ], - } - ), - ) as mock_request: - # Mock logging object + mock_http_response = MagicMock() + mock_http_response.status_code = 200 + mock_http_response.json.return_value = { + "created": 1234567890, + "data": [{"url": "https://example.com/image.png", "revised_prompt": prompt}], + } + + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_http_response) + + with patch( + "litellm.llms.azure.azure.get_async_httpx_client", return_value=mock_client + ): logging_obj = MagicMock() logging_obj.pre_call = MagicMock() logging_obj.post_call = MagicMock() - # Call the aimage_generation method - response = await azure_chat_completion.aimage_generation( + await azure_chat_completion.aimage_generation( data=data, model_response=None, azure_client_params=azure_client_params, @@ -430,30 +410,14 @@ async def test_azure_aimage_generation_base_model_vs_deployment_name(): input=[], logging_obj=logging_obj, headers={}, - model=model, # Pass the deployment name + model=model, timeout=60.0, ) - # Verify the mock was called - assert mock_request.called, "HTTP request should have been made" - - # Get the call arguments - call_kwargs = mock_request.call_args.kwargs - - # Verify the URL uses the deployment name (not base_model) - api_base_used = call_kwargs.get("api_base", "") - assert model in api_base_used, ( - f"URL should contain deployment name '{model}', " - f"but got: {api_base_used}" - ) - assert base_model not in api_base_used or base_model == model, ( - f"URL should NOT contain base_model '{base_model}' when it differs from deployment name, " - f"but got: {api_base_used}" - ) - - request_data = call_kwargs.get("data", {}) - wire_json = AzureChatCompletion.azure_deployment_image_generation_json_body( - api_base_used, request_data - ) + assert mock_client.post.called + post_kwargs = mock_client.post.call_args.kwargs + url_used = post_kwargs.get("url", "") + assert model in url_used + wire_json = post_kwargs.get("json") or {} assert "model" not in wire_json - assert request_data.get("model") == base_model + assert data.get("model") == base_model