diff --git a/litellm/images/main.py b/litellm/images/main.py index 03c0e36ad93..cf588cbcf0f 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -2,7 +2,18 @@ import asyncio import contextvars import importlib from functools import partial -from typing import TYPE_CHECKING, Any, Coroutine, Dict, List, Literal, Optional, Union, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + Coroutine, + Dict, + List, + Literal, + Optional, + Union, + cast, + overload, +) if TYPE_CHECKING: from litellm.images.utils import ImageEditRequestUtils @@ -10,7 +21,7 @@ if TYPE_CHECKING: import httpx import litellm -from litellm.utils import exception_type, get_litellm_params + # client is imported from litellm as it's a decorator from litellm import client from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL @@ -23,6 +34,7 @@ from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.custom_llm import CustomLLM +from litellm.utils import exception_type, get_litellm_params #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -32,8 +44,8 @@ from litellm.main import ( azure_chat_completions, base_llm_aiohttp_handler, base_llm_http_handler, - bedrock_image_generation, bedrock_image_edit, + bedrock_image_generation, openai_chat_completions, openai_image_variations, ) @@ -330,11 +342,36 @@ def image_generation( # noqa: PLR0915 azure_ad_token = optional_params.pop( "azure_ad_token", None ) or get_secret_str("AZURE_AD_TOKEN") + + # Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided + if azure_ad_token_provider is None: + from litellm.llms.azure.common_utils import ( + get_azure_ad_token_from_entra_id, + ) + + # Extract Azure AD credentials from litellm_params + tenant_id = litellm_params_dict.get("tenant_id") + client_id = litellm_params_dict.get("client_id") + client_secret = litellm_params_dict.get("client_secret") + azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default" + + # Create token provider if credentials are available + if tenant_id and client_id and client_secret: + azure_ad_token_provider = get_azure_ad_token_from_entra_id( + tenant_id=tenant_id, + client_id=client_id, + client_secret=client_secret, + scope=azure_scope, + ) default_headers = { "Content-Type": "application/json", - "api-key": api_key, } + # Only add api-key header if api_key is not None + # Azure AD authentication will use Authorization header instead + if api_key is not None: + default_headers["api-key"] = api_key + for k, v in default_headers.items(): if k not in headers: headers[k] = v @@ -399,8 +436,12 @@ def image_generation( # noqa: PLR0915 default_headers = { "Content-Type": "application/json", - "api-key": api_key, } + # Only add api-key header if api_key is not None + # Azure AD authentication will use Authorization header instead + if api_key is not None: + default_headers["api-key"] = api_key + for k, v in default_headers.items(): if k not in headers: headers[k] = v @@ -983,6 +1024,7 @@ def __getattr__(name: str) -> Any: if name == "ImageEditRequestUtils": # Lazy load ImageEditRequestUtils to avoid heavy import from images.utils at module load time from .utils import ImageEditRequestUtils as _ImageEditRequestUtils + # Cache it in the module's __dict__ for subsequent accesses module = importlib.import_module(__name__) module.__dict__["ImageEditRequestUtils"] = _ImageEditRequestUtils 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 f31001ebd36..998510efcd9 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 @@ -3,7 +3,7 @@ import os import sys import traceback from typing import Callable, Optional -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -87,3 +87,80 @@ def test_azure_image_generation_flattens_extra_body(): assert data["custom_param"] == "test_value" assert data["n"] == 1 assert data["size"] == "1024x1024" + + +def test_azure_image_generation_creates_token_provider_from_credentials(): + """ + Test that azure_ad_token_provider is created from tenant_id, client_id, client_secret. + + This test verifies the fix in images/main.py where we now create the + azure_ad_token_provider from credentials in litellm_params if it's not already provided. + """ + # Simulate the fix in images/main.py + litellm_params_dict = { + "tenant_id": "test-tenant-id", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + "azure_scope": None, + } + + azure_ad_token_provider = None + + # This is the logic we added in images/main.py + if azure_ad_token_provider is None: + tenant_id = litellm_params_dict.get("tenant_id") + client_id = litellm_params_dict.get("client_id") + client_secret = litellm_params_dict.get("client_secret") + azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default" + + # Verify the credentials are extracted correctly + assert tenant_id == "test-tenant-id" + assert client_id == "test-client-id" + assert client_secret == "test-client-secret" + assert azure_scope == "https://cognitiveservices.azure.com/.default" + + # Verify the condition to create token provider is met + assert tenant_id and client_id and client_secret, "Credentials should be present to create token provider" + + +def test_azure_image_generation_headers_without_api_key(): + """ + Test that when api_key is None, the api-key header is not added to headers. + + This prevents the httpx TypeError: "Header value must be str or bytes, not " + that was occurring when api_key was None and being set in headers. + + This is a unit test for the fix in images/main.py where we now check: + if api_key is not None: + default_headers["api-key"] = api_key + """ + from litellm.images.main import image_generation + + # Test the header building logic directly + api_key = None + + default_headers = { + "Content-Type": "application/json", + } + + # This is the fix: only add api-key if it's not None + if api_key is not None: + default_headers["api-key"] = api_key + + # Verify api-key is not in headers when api_key is None + assert "api-key" not in default_headers + + # Verify Content-Type is still there + assert default_headers["Content-Type"] == "application/json" + + # Test with a valid api_key + api_key = "valid-key-123" + default_headers_with_key = { + "Content-Type": "application/json", + } + if api_key is not None: + default_headers_with_key["api-key"] = api_key + + # Verify api-key is added when api_key is valid + assert "api-key" in default_headers_with_key + assert default_headers_with_key["api-key"] == "valid-key-123"