diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 4fc1ae960b..e1ac185891 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,3 +1,5 @@ +import asyncio +import hashlib import json import os from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast @@ -449,6 +451,25 @@ class BaseAzureLLM(BaseOpenAILLM): ] = None client_initialization_params: dict = locals() client_initialization_params["is_async"] = _is_async + _lp = litellm_params or {} + _ad_provider = _lp.get("azure_ad_token_provider") + _ad_token = _lp.get("azure_ad_token") + _client_secret = _lp.get("client_secret") + _azure_password = _lp.get("azure_password") + client_initialization_params["azure_ad_token"] = ( + hashlib.sha256(_ad_token.encode()).hexdigest() + if isinstance(_ad_token, str) + else None + ) + client_initialization_params["azure_ad_token_provider"] = ( + f"provider_id={id(_ad_provider) if callable(_ad_provider) else None}" + f"|tenant_id={_lp.get('tenant_id')}" + f"|client_id={_lp.get('client_id')}" + f"|client_secret={hashlib.sha256(_client_secret.encode()).hexdigest() if isinstance(_client_secret, str) else None}" + f"|azure_username={_lp.get('azure_username')}" + f"|azure_password={hashlib.sha256(_azure_password.encode()).hexdigest() if isinstance(_azure_password, str) else None}" + f"|azure_scope={_lp.get('azure_scope')}" + ) if client is None: cached_client = self.get_cached_openai_client( client_initialization_params=client_initialization_params, @@ -474,8 +495,29 @@ class BaseAzureLLM(BaseOpenAILLM): if self._is_azure_v1_api_version(api_version): # Extract only params that OpenAI client accepts # Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview" - v1_params = { - "api_key": azure_client_params.get("api_key"), + # The OpenAI client accepts a callable for `api_key` and re-invokes it + # on every request (via `_refresh_api_key`), so passing + # `azure_ad_token_provider` directly preserves Azure AD token refresh + # behavior that the regular AzureOpenAI client provides. + v1_api_key: Optional[Union[str, Callable[[], Any]]] = ( + azure_client_params.get("api_key") + or azure_client_params.get("azure_ad_token_provider") + or azure_client_params.get("azure_ad_token") + ) + if _is_async is True and callable(v1_api_key): + # AsyncOpenAI expects an async provider; wrap the sync provider + # returned by azure-identity. Offload to a thread so a token + # refresh (blocking HTTP call to AAD on cache miss) does not + # stall the event loop. + _sync_provider = v1_api_key + + async def _async_v1_api_key() -> str: + return await asyncio.to_thread(_sync_provider) + + v1_api_key = _async_v1_api_key + + v1_params: Dict[str, Any] = { + "api_key": v1_api_key, "base_url": f"{api_base}/openai/v1/", } if "timeout" in azure_client_params: diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 3fa794375e..413241adf3 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -1646,6 +1646,336 @@ def test_azure_v1_api_uses_openai_client(api_version): ), f"base_url should contain /openai/v1/, got {async_client.base_url}" +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_with_azure_ad_token_provider(api_version): + """ + The v1 OpenAI client path must forward `azure_ad_token_provider` so Azure AD + auth works for `api_version` in {"v1", "latest", "preview"}. + + Regression: https://github.com/BerriAI/litellm/issues/27945 — before the fix + the v1 branch only forwarded `api_key`, so AD-only configs raised + "The api_key client option must be set" on every request. + + The OpenAI SDK accepts a callable for `api_key` and re-invokes it on every + request, so passing the provider directly preserves token refresh. + """ + from openai import AsyncOpenAI, OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + token_value = "mock-azure-ad-token-from-provider" + + def token_provider(): + return token_value + + init_return = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": token_provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + + client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + # The SDK stores callables as `_api_key_provider` and refreshes + # `self.api_key` before each request. + assert client._api_key_provider is token_provider + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + + async_client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + assert isinstance(async_client, AsyncOpenAI) + # Async client requires an async provider; we wrap the sync provider + # so the SDK can `await` it. + assert async_client._api_key_provider is not None + assert async_client._api_key_provider is not token_provider + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_async_token_provider_resolves_to_current_token(api_version): + """ + The async wrapper must call the underlying sync provider on each invocation + (not cache its first return value), so token rotation is honored. + """ + import asyncio + + from openai import AsyncOpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + tokens = iter(["token-1", "token-2", "token-3"]) + + def rotating_provider(): + return next(tokens) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": rotating_provider, + } + + async_client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + assert isinstance(async_client, AsyncOpenAI) + loop = asyncio.new_event_loop() + try: + first = loop.run_until_complete(async_client._api_key_provider()) + second = loop.run_until_complete(async_client._api_key_provider()) + finally: + loop.close() + + assert first == "token-1" + assert second == "token-2" + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_with_static_azure_ad_token(api_version): + """ + When only `azure_ad_token` (a static string) is set, the v1 client should + receive it as `api_key`. + """ + from openai import OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + token_value = "static-azure-ad-token" + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": token_value, + "azure_ad_token_provider": None, + } + + client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + assert client.api_key == token_value + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_key_wins_over_ad_token(api_version): + """ + Explicit `api_key` takes precedence over `azure_ad_token_provider` / + `azure_ad_token`, matching the priority documented in + `initialize_azure_sdk_client`. + """ + from openai import OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "explicit-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": "should-be-ignored", + "azure_ad_token_provider": lambda: "also-ignored", + } + + client = base_llm.get_azure_openai_client( + api_key="explicit-key", + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + assert client.api_key == "explicit-key" + assert client._api_key_provider is None + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_separates_distinct_ad_providers(api_version): + """ + Two configs sharing api_base/api_version but with different AD token + providers must not share a cached OpenAI client, otherwise requests for + one config would be sent with another config's AD credentials. + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def provider_a(): + return "token-a" + + def provider_b(): + return "token-b" + + def _init_for(provider): + return { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_for(provider_a) + client_a = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider_a}, + _is_async=True, + ) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_for(provider_b) + client_b = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider_b}, + _is_async=True, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert isinstance(client_b, AsyncOpenAI) + assert client_a is not client_b + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_separates_distinct_entra_credentials(api_version): + """ + Configs that synthesize an AD provider from tenant_id/client_id/client_secret + must not share a cached client when those inputs differ. + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def synth_provider(): + return "synthesized-token" + + def _init_synth(): + return { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": synth_provider, + } + + common = { + "api_key": None, + "api_base": api_base, + "api_version": api_version, + "_is_async": True, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_synth() + client_a = base_llm.get_azure_openai_client( + litellm_params={ + "tenant_id": "tenant-a", + "client_id": "client-a", + "client_secret": "secret-a", + }, + **common, + ) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_synth() + client_b = base_llm.get_azure_openai_client( + litellm_params={ + "tenant_id": "tenant-b", + "client_id": "client-b", + "client_secret": "secret-b", + }, + **common, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert isinstance(client_b, AsyncOpenAI) + assert client_a is not client_b + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_reuses_for_identical_ad_config(api_version): + """ + Identical AD configs should still share a cached client (regression guard + so the cache-key change doesn't accidentally disable caching). + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def provider(): + return "tok" + + init_return = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + client_a = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider}, + _is_async=True, + ) + client_b = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider}, + _is_async=True, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert client_a is client_b + + def test_azure_traditional_api_uses_azure_openai_client(): """ Test that traditional Azure API versions still use AzureOpenAI client.