From c008f25b55cc75da65ec47379c8b2ac1f59d3eb6 Mon Sep 17 00:00:00 2001 From: mateo Date: Mon, 19 Oct 2026 16:20:20 +0000 Subject: [PATCH] fix(azure): cache the remaining per-request Entra credentials and keep token refresh off the event loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/azure/common_utils.py | 97 ++++++++++++++----- litellm/llms/azure/search/transformation.py | 5 +- litellm/llms/custom_httpx/llm_http_handler.py | 33 ++++--- litellm/realtime_api/main.py | 23 +++-- litellm/secret_managers/main.py | 9 +- ...st_bing_grounding_search_transformation.py | 21 ++++ .../llms/azure/test_azure_common_utils.py | 78 ++++++++++++++- .../custom_httpx/test_llm_http_handler.py | 50 ++++++++++ 8 files changed, 268 insertions(+), 48 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index f9f636244f4..fdf19eb9cf8 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -3,13 +3,14 @@ import hashlib import json import os import re -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from functools import lru_cache from types import MappingProxyType from typing import Final, Literal, NamedTuple, cast import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -38,6 +39,16 @@ class _AzureV1ClientParams(TypedDict, total=False, extra_items=object): base_url: ReadOnly[str] +class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object): + api_key: ReadOnly[str | None] + azure_ad_token: ReadOnly[str | None] + azure_ad_token_provider: ReadOnly[Callable[[], str | Awaitable[str]] | None] + + +_ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams) +_AZURE_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter(Callable[[], str]) + + class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object): api_version: ReadOnly[str] base_url: ReadOnly[str] @@ -98,11 +109,41 @@ def _cached_entra_id_token_provider( return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) +@lru_cache(maxsize=128) +def _cached_username_password_token_provider( + client_id: str, + azure_username: str, + azure_password: str, + scope: str, +) -> Callable[[], str]: + from azure.identity import UsernamePasswordCredential, get_bearer_token_provider + + credential: Final = UsernamePasswordCredential( + client_id=client_id, + username=azure_username, + password=azure_password, + ) + return get_bearer_token_provider(credential, scope) + + @lru_cache(maxsize=128) def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=scope) +def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: + return _cached_azure_ad_token_refresh_provider(scope) + + +def _async_azure_ad_token_provider( + sync_provider: Callable[[], str], +) -> Callable[[], Awaitable[str]]: + async def async_provider() -> str: + return await asyncio.to_thread(sync_provider) + + return async_provider + + def get_azure_ad_token_from_entra_id( tenant_id: str, client_id: str, @@ -177,22 +218,19 @@ def get_azure_ad_token_from_username_password( Returns: callable that returns a bearer token. """ - from azure.identity import UsernamePasswordCredential, get_bearer_token_provider - verbose_logger.debug( "client_id=%s, azure_username=[set=%s], azure_password=[set=%s]", client_id, azure_username is not None, azure_password is not None, ) - credential: Final = UsernamePasswordCredential( + token_provider: Final = _cached_username_password_token_provider( client_id=client_id, - username=azure_username, - password=azure_password, + azure_username=azure_username, + azure_password=azure_password, + scope=scope, ) - token_provider: Final = get_bearer_token_provider(credential, scope) - verbose_logger.debug("token_provider %s", token_provider) return token_provider @@ -560,22 +598,19 @@ class BaseAzureLLM(BaseOpenAILLM): # 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: str | Callable[[], object] | None = ( - azure_client_params.get("api_key") - or azure_client_params.get("azure_ad_token_provider") - or azure_client_params.get("azure_ad_token") + v1_client_params: Final[_AsyncAzureClientParams] = _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER.validate_python( + azure_client_params + ) + v1_token_provider: Final = ( + v1_client_params.get("api_key") + or v1_client_params.get("azure_ad_token_provider") + or v1_client_params.get("azure_ad_token") + ) + v1_api_key: Final[str | Callable[[], str | Awaitable[str]] | None] = ( + _async_azure_ad_token_provider(_AZURE_TOKEN_PROVIDER_ADAPTER.validate_python(v1_token_provider)) + if _is_async is True and callable(v1_token_provider) + else v1_token_provider ) - 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: Final = 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: Final[_AzureV1ClientParams] = { "api_key": v1_api_key, @@ -597,7 +632,21 @@ class BaseAzureLLM(BaseOpenAILLM): else: # Traditional Azure API uses AzureOpenAI client if _is_async is True: - openai_client = AsyncAzureOpenAI(**azure_client_params) + async_azure_client_params: Final[_AsyncAzureClientParams] = ( + _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER.validate_python(azure_client_params) + ) + async_token_provider_candidate: Final = async_azure_client_params.get("azure_ad_token_provider") + async_azure_client_params_with_threaded_provider: Final[_AsyncAzureClientParams] = ( + { + **async_azure_client_params, + "azure_ad_token_provider": _async_azure_ad_token_provider( + _AZURE_TOKEN_PROVIDER_ADAPTER.validate_python(async_token_provider_candidate) + ), + } + if callable(async_token_provider_candidate) + else async_azure_client_params + ) + openai_client = AsyncAzureOpenAI(**async_azure_client_params_with_threaded_provider) else: openai_client = AzureOpenAI(**azure_client_params) diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index 4264dae018a..f2bb2ecc68c 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -35,6 +35,7 @@ from typing import TYPE_CHECKING, Final, Literal import httpx from pydantic import ConfigDict, ValidationError +from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, @@ -263,9 +264,7 @@ def _search_tool(optional_params: Mapping[str, object]) -> _BingGroundingTool | def _default_entra_token_minter() -> str: - from litellm.secret_managers.get_azure_ad_token_provider import get_azure_ad_token_provider - - return get_azure_ad_token_provider(azure_scope=ENTRA_SCOPE)() + return cached_azure_ad_token_refresh_provider(ENTRA_SCOPE)() class BingGroundingSearchConfig(BaseSearchConfig): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b03b8050299..66ed80b7afb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2953,14 +2953,15 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers = responses_api_provider_config.validate_environment( + validated_headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers=response_api_optional_request_params.get("extra_headers", {}) or {}, model=model, litellm_params=litellm_params, ) if extra_headers: - headers.update(extra_headers) + validated_headers.update(extra_headers) # Check if streaming is requested stream = response_api_optional_request_params.get("stream", False) @@ -2976,7 +2977,7 @@ class BaseLLMHTTPHandler: input=input, response_api_optional_request_params=response_api_optional_request_params, litellm_params=litellm_params, - headers=headers, + headers=validated_headers, ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) @@ -3011,7 +3012,7 @@ class BaseLLMHTTPHandler: headers, signed_body = await sign_request_off_loop_if_aws( responses_api_provider_config, responses_api_provider_config.sign_request, - headers=headers, + headers=validated_headers, optional_params=dict(litellm_params), request_data=data, api_base=api_base, @@ -3145,7 +3146,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = responses_api_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers=extra_headers or {}, model="None", litellm_params=litellm_params ) @@ -3399,7 +3401,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = responses_api_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers=extra_headers or {}, model="None", litellm_params=litellm_params ) @@ -3563,7 +3566,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = responses_api_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers=extra_headers or {}, model="None", litellm_params=litellm_params ) @@ -4620,7 +4624,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = responses_api_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers=extra_headers or {}, model="None", litellm_params=litellm_params ) @@ -4795,12 +4800,13 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers = responses_api_provider_config.validate_environment( + validated_headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers=extra_headers or {}, model=model, litellm_params=litellm_params ) if extra_headers: - headers.update(extra_headers) + validated_headers.update(extra_headers) api_base: Final = responses_api_provider_config.get_complete_url( api_base=litellm_params.api_base, @@ -4816,14 +4822,14 @@ class BaseLLMHTTPHandler: response_api_optional_request_params=response_api_optional_request_params, api_base=api_base, litellm_params=litellm_params, - headers=headers, + headers=validated_headers, ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) headers, signed_body = await sign_request_off_loop_if_aws( responses_api_provider_config, responses_api_provider_config.sign_request, - headers=headers, + headers=validated_headers, optional_params=dict(litellm_params), request_data=data, api_base=url, @@ -6991,7 +6997,8 @@ class BaseLLMHTTPHandler: **kwargs, } ) - headers: Final = responses_api_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + responses_api_provider_config.validate_environment, headers={}, model=model, litellm_params=litellm_params, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 96b4b081c05..ea47d744cae 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -562,7 +562,12 @@ async def _arealtime( configured_realtime_protocol, query_params=query_params, websocket=websocket ) resolved_azure_ad_token: Final = ( - None if api_key else get_azure_ad_token(GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token)) + None + if api_key + else await asyncio.to_thread( + get_azure_ad_token, + GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token), + ) ) await azure_realtime.async_realtime( model=model, @@ -743,15 +748,21 @@ def _azure_realtime_health_protocol( return "GA", query_params -def _realtime_health_check_auth_headers( +async def _realtime_health_check_auth_headers( custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, object] ) -> Mapping[str, str]: if custom_llm_provider == "azure": + azure_ad_token: Final = ( + None + if api_key + else await asyncio.to_thread( + get_azure_ad_token, + GenericLiteLLMParams.model_validate(dict(model_params)), + ) + ) return azure_realtime.get_auth_headers( api_key=api_key, - azure_ad_token=( - None if api_key else get_azure_ad_token(GenericLiteLLMParams.model_validate(dict(model_params))) - ), + azure_ad_token=azure_ad_token, ) if api_key is None: return _EMPTY_AUTH_HEADERS @@ -807,7 +818,7 @@ async def realtime_health_check( if custom_llm_provider == LlmProviders.OPENAI.value else resolved_api_key ) - auth_headers: Final = _realtime_health_check_auth_headers( + auth_headers: Final = await _realtime_health_check_auth_headers( custom_llm_provider=custom_llm_provider, api_key=health_check_api_key, model_params=resolved_params, diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 716407d6def..d12e200772f 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -3,6 +3,8 @@ import base64 import os import time import traceback +from collections.abc import Callable +from functools import lru_cache from typing import Final import httpx @@ -23,6 +25,11 @@ oidc_cache: Final = DualCache() _PARSED_LITERAL: Final = TypeAdapter(object) +@lru_cache(maxsize=128) +def _cached_azure_oidc_token_provider(audience: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=audience) + + _OIDC_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60 @@ -283,7 +290,7 @@ def get_secret( verbose_logger.warning( "AZURE_FEDERATED_TOKEN_FILE not found in environment will use Azure AD token provider" ) - azure_token_provider: Final = get_azure_ad_token_provider(azure_scope=oidc_aud) + azure_token_provider: Final = _cached_azure_oidc_token_provider(oidc_aud) try: oidc_token = azure_token_provider() if oidc_token is None: diff --git a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py index f8d042c14fd..5a97b6f933c 100644 --- a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py +++ b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py @@ -110,6 +110,27 @@ def test_validate_environment_entra_minter_failure_names_the_options(): assert "no az login" in message +def test_default_entra_minter_reuses_provider(monkeypatch: pytest.MonkeyPatch): + import importlib + + from litellm.llms.azure import common_utils + from litellm.llms.azure.search.transformation import _default_entra_token_minter + + provider_module = importlib.import_module("litellm.secret_managers.get_azure_ad_token_provider") + provider_factory = Mock(side_effect=lambda **_: lambda: "entra-token") + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + monkeypatch.setattr(common_utils, "get_azure_ad_token_provider", provider_factory) + monkeypatch.setattr(provider_module, "get_azure_ad_token_provider", provider_factory) + + try: + first = _default_entra_token_minter() + second = _default_entra_token_minter() + assert first == second == "entra-token" + assert provider_factory.call_count == 1 + finally: + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + def test_validate_environment_does_not_mutate_and_is_idempotent(): config = _config() caller_headers = {"X-Custom": "keep-me"} diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 80449909b3f..4a4201c5fe9 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1,7 +1,10 @@ +import asyncio +import itertools import json import os +import time import traceback -from typing import Callable, Optional +from typing import Callable, Final, Optional from unittest.mock import MagicMock, patch import pytest @@ -13,6 +16,7 @@ from litellm.llms.azure.common_utils import ( _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, + get_azure_ad_token_from_username_password, ) from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -2141,6 +2145,78 @@ class TestEntraIdTokenProviderCache: assert mock_credential.call_count == 2 +def test_username_password_token_provider_reuses_same_credentials_and_separates_passwords(): + with ( + patch("azure.identity.UsernamePasswordCredential") as mock_credential, + patch( + "azure.identity.get_bearer_token_provider", + side_effect=lambda credential, scope: lambda: "token", + ), + ): + first = get_azure_ad_token_from_username_password( + client_id="client", + azure_username="user", + azure_password="password", + scope="https://cognitiveservices.azure.com/.default", + ) + second = get_azure_ad_token_from_username_password( + client_id="client", + azure_username="user", + azure_password="password", + scope="https://cognitiveservices.azure.com/.default", + ) + changed_password = get_azure_ad_token_from_username_password( + client_id="client", + azure_username="user", + azure_password="different-password", + scope="https://cognitiveservices.azure.com/.default", + ) + + assert first() == second() == changed_password() == "token" + assert mock_credential.call_count == 2 + + +@pytest.mark.asyncio +async def test_non_v1_azure_token_refresh_does_not_block_the_event_loop(): + from openai import AsyncAzureOpenAI + + base_llm = BaseAzureLLM() + api_base: Final = "https://test.openai.azure.com" + + def slow_token_provider() -> str: + time.sleep(0.3) + return "slow-token" + + client: Final = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version="2024-02-01", + litellm_params={"azure_ad_token_provider": slow_token_provider}, + _is_async=True, + ) + + assert isinstance(client, AsyncAzureOpenAI) + ticks: Final = itertools.count() + token_resolved: Final = asyncio.Event() + + async def resolve_token() -> None: + try: + assert await client._get_azure_ad_token() == "slow-token" + finally: + token_resolved.set() + + async def ticker() -> None: + while not token_resolved.is_set(): + next(ticks) + await asyncio.sleep(0.01) + + try: + await asyncio.gather(resolve_token(), ticker()) + assert next(ticks) > 0 + finally: + await client.close() + + def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client. diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 23c142b61ab..0d4825152cb 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1,6 +1,7 @@ import asyncio import base64 import inspect +import itertools import json import logging import threading @@ -476,6 +477,55 @@ async def test_async_response_api_handler_streams_when_provider_transform_adds_s assert client.post.call_args.kwargs["json"]["stream"] is True +@pytest.mark.asyncio +async def test_async_response_api_handler_validates_environment_off_the_event_loop(): + validation_started: Final = threading.Event() + validation_finished: Final = threading.Event() + ticks: Final = itertools.count() + + def validate_environment(**_kwargs: object) -> dict[str, str]: + validation_started.set() + time.sleep(0.3) + validation_finished.set() + return {} + + config = Mock() + config.validate_environment.side_effect = validate_environment + config.get_complete_url.return_value = "https://responses.example.com/v1/responses" + config.async_transform_responses_api_request = AsyncMock(return_value={"model": "model", "input": "hi"}) + config.sign_request.return_value = ({}, None) + client = AsyncMock(spec=AsyncHTTPHandler) + client.post.return_value = httpx.Response( + 200, + request=httpx.Request("POST", "https://responses.example.com/v1/responses"), + ) + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = None + + async def call_handler() -> None: + await BaseLLMHTTPHandler().async_response_api_handler( + model="model", + input="hi", + responses_api_provider_config=config, + response_api_optional_request_params={}, + custom_llm_provider="test", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + client=client, + ) + + async def ticker() -> None: + while not validation_finished.is_set(): + if validation_started.is_set(): + next(ticks) + await asyncio.sleep(0.01) + + await asyncio.gather(call_handler(), ticker()) + + assert next(ticks) > 0 + client.post.assert_awaited_once() + + _CHATGPT_SSE_BODY = ( "event: response.output_item.done\n" 'data: {"type": "response.output_item.done", "output_index": 0, "item": {"type": "message", '