From a1b254f263d539bce98e8760439d28df61c63e82 Mon Sep 17 00:00:00 2001 From: Frederic Kayser Date: Wed, 29 Apr 2026 10:29:19 +0100 Subject: [PATCH] Implement Azure AD token authentication and update header checks in image edit tests --- .../llms/azure/image_edit/transformation.py | 20 +-- tests/image_gen_tests/test_image_edits.py | 127 +++++++++++++++++- 2 files changed, 129 insertions(+), 18 deletions(-) diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index dffa1c9eea5..6f9abdf7ff6 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -3,8 +3,10 @@ from typing import Optional, cast import httpx import litellm +from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams from litellm.utils import _add_path_to_api_base @@ -17,21 +19,13 @@ class AzureImageEditConfig(OpenAIImageEditConfig): litellm_params: Optional[dict] = None, api_base: 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") + _litellm_params = GenericLiteLLMParams(**(litellm_params or {})) + if api_key: + _litellm_params.api_key = api_key + return BaseAzureLLM._base_validate_azure_environment( + headers=headers, litellm_params=_litellm_params ) - headers.update( - { - "Authorization": f"Bearer {api_key}", - } - ) - return headers - def get_complete_url( self, model: str, diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 3504065a109..1fc8da88772 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -326,13 +326,13 @@ async def test_azure_image_edit_litellm_sdk(): prompt.strip() in form_data["prompt"] ), f"Expected prompt to contain '{prompt.strip()}'" - # Check headers + # Check headers - API key auth uses the api-key header 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" + assert "api-key" in headers, "api-key header should be present for API key auth" + assert ( + headers["api-key"] == test_api_key + ), f"Expected api-key '{test_api_key}', got {headers['api-key']}" print("result from image edit", result) @@ -780,3 +780,120 @@ async def test_image_edit_array_handling(): # Verify that both calls were made to the API assert mock_post.call_count == 2 + + +@pytest.mark.asyncio +async def test_azure_image_edit_azure_ad_token_auth(): + """Test Azure image edit uses Azure AD token when no API key is provided.""" + from litellm import aimage_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 + + fake_ad_token = "fake-azure-ad-token-12345" + token_provider = lambda: fake_ad_token + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(mock_response, 200) + + test_api_base = "https://ai-api-gw-uae-north.openai.azure.com" + test_api_version = "2025-04-01-preview" + + result = await aimage_edit( + prompt="Edit this image", + model="azure/gpt-image-1", + api_base=test_api_base, + api_version=test_api_version, + azure_ad_token_provider=token_provider, + image=TEST_IMAGES, + ) + + mock_post.assert_called_once() + call_args = mock_post.call_args + headers = call_args.kwargs.get("headers", {}) + + # Azure AD token auth should use Authorization: Bearer header + assert "Authorization" in headers, "Authorization header should be present for Azure AD auth" + assert headers["Authorization"] == f"Bearer {fake_ad_token}", ( + f"Expected Bearer token with AD token, got {headers['Authorization']}" + ) + # api-key header should NOT be present + assert "api-key" not in headers, "api-key header should not be present for Azure AD auth" + + ImageResponse.model_validate(result) + + +@pytest.mark.asyncio +async def test_azure_image_edit_api_key_takes_precedence(): + """Test that API key auth takes precedence over Azure AD token.""" + from litellm import aimage_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 + + fake_ad_token = "fake-azure-ad-token-12345" + token_provider = lambda: fake_ad_token + test_api_key = "test-api-key-priority" + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(mock_response, 200) + + result = await aimage_edit( + prompt="Edit this image", + model="azure/gpt-image-1", + api_base="https://test.openai.azure.com", + api_key=test_api_key, + api_version="2025-04-01-preview", + azure_ad_token_provider=token_provider, + image=TEST_IMAGES, + ) + + mock_post.assert_called_once() + call_args = mock_post.call_args + headers = call_args.kwargs.get("headers", {}) + + # API key should take precedence — uses api-key header + assert "api-key" in headers, "api-key header should be present when API key is provided" + assert headers["api-key"] == test_api_key + # Authorization header should NOT be present + assert "Authorization" not in headers, ( + "Authorization header should not be set when API key is available" + ) + + ImageResponse.model_validate(result)