From d35691aa0ce3e40bb4555fed715d9e45ec64e361 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 9 Feb 2026 16:25:46 +0530 Subject: [PATCH] Fix: base_model name for body and deplyment name in URL --- litellm/llms/azure/azure.py | 18 +- .../test_azure_image_generation_init.py | 205 ++++++++++++++++++ 2 files changed, 214 insertions(+), 9 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index cb9fe0aeb30..95afa53b625 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1060,6 +1060,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, client=None, timeout=None, + model: Optional[str] = None, ) -> ImageResponse: response: Optional[dict] = None @@ -1071,8 +1072,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): 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=data.get("model", "") + azure_client_params=azure_client_params, model=model or data.get("model", "") ) ## LOGGING @@ -1159,21 +1161,18 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): model = model else: model = None - ## BASE MODEL CHECK if ( model_response is not None - and optional_params.get("base_model", None) is not None + and litellm_params.get("base_model", None) is not None ): - model_response._hidden_params["model"] = optional_params.pop( - "base_model" - ) + model_response._hidden_params["model"] = litellm_params.get("base_model", None) # Azure image generation API doesn't support extra_body parameter extra_body = optional_params.pop("extra_body", {}) flattened_params = {**optional_params, **extra_body} - data = {"model": model, "prompt": prompt, **flattened_params} + data = {"model": litellm_params.get("base_model", None) or model, "prompt": prompt, **flattened_params} max_retries = data.pop("max_retries", 2) if not isinstance(max_retries, int): raise AzureOpenAIError( @@ -1196,10 +1195,11 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): is_async=False, ) 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) # type: ignore + 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=data.get("model", "") + azure_client_params=azure_client_params, model=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 index 987eb5bf998..bf8165c6908 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 @@ -251,3 +251,208 @@ def test_azure_image_generation_drop_params_false_raises_error(): # Verify the error message mentions the unsupported parameter assert "response_format" in str(exc_info.value) + + +def test_azure_image_generation_base_model_vs_deployment_name(): + """ + Test that Azure image generation correctly uses base_model in request body + but deployment name in the URL. + + When base_model is specified in litellm_params, the request should: + 1. Use base_model (e.g., "gpt-image-1.5") in the JSON request body + 2. Use the deployment name (e.g., "gpt-image-15") in the URL path + + This is important because Azure expects: + - URL: /openai/deployments/{deployment_name}/images/generations + - Body: {"model": "{base_model}", ...} + + Example config: + model: azure/gpt-image-15 # deployment name + base_model: gpt-image-1.5 # actual model name + """ + from unittest.mock import MagicMock + + # Setup test parameters + azure_chat_completion = AzureChatCompletion() + + prompt = "A beautiful image of a cat" + model = "gpt-image-15" # This is the deployment name + base_model = "gpt-image-1.5" # This is the actual model name + api_base = "https://openai-gpt-image-1-5-test-v-1.openai.azure.com/" + api_version = "2024-07-01-preview" + api_key = "test-api-key" + + litellm_params = { + "base_model": base_model, + "api_base": api_base, + "api_version": api_version, + } + + optional_params = { + "n": 1, + "size": "1024x1024" + } + + # Mock the HTTP request to capture what gets sent + 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 + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.post_call = MagicMock() + + # Call the image_generation method + try: + response = azure_chat_completion.image_generation( + prompt=prompt, + timeout=60.0, + optional_params=optional_params, + logging_obj=logging_obj, + headers={}, + model=model, + api_key=api_key, + api_base=api_base, + api_version=api_version, + litellm_params=litellm_params, + ) + except Exception as e: + # If there's an error, we still want to check the mock calls + pass + + # 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}" + ) + + # Verify the request body uses base_model (not deployment name) + request_data = call_kwargs.get("data", {}) + assert request_data.get("model") == base_model, ( + f"Request body 'model' field should be base_model '{base_model}', " + f"but got: {request_data.get('model')}" + ) + + # Verify other fields are correct + assert request_data.get("prompt") == prompt + assert request_data.get("n") == 1 + assert request_data.get("size") == "1024x1024" + + +@pytest.mark.asyncio +async def test_azure_aimage_generation_base_model_vs_deployment_name(): + """ + Test that Azure async image generation correctly uses base_model in request body + but deployment name in the URL. + + This is the async version of test_azure_image_generation_base_model_vs_deployment_name. + """ + from unittest.mock import MagicMock + + # Setup test parameters + azure_chat_completion = AzureChatCompletion() + + prompt = "A beautiful image of a cat" + model = "gpt-image-15" # This is the deployment name + base_model = "gpt-image-1.5" # This is the actual model name + api_base = "https://openai-gpt-image-1-5-test-v-1.openai.azure.com/" + api_version = "2024-07-01-preview" + api_key = "test-api-key" + + data = { + "model": base_model, + "prompt": prompt, + "n": 1, + "size": "1024x1024" + } + + azure_client_params = { + "api_base": api_base, + "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 + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.post_call = MagicMock() + + # Call the aimage_generation method + try: + response = await azure_chat_completion.aimage_generation( + data=data, + model_response=None, + azure_client_params=azure_client_params, + api_key=api_key, + input=[], + logging_obj=logging_obj, + headers={}, + model=model, # Pass the deployment name + timeout=60.0, + ) + except Exception as e: + # If there's an error, we still want to check the mock calls + pass + + # 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}" + ) + + # Verify the request body uses base_model (not deployment name) + request_data = call_kwargs.get("data", {}) + assert request_data.get("model") == base_model, ( + f"Request body 'model' field should be base_model '{base_model}', " + f"but got: {request_data.get('model')}" + )