scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 199

fix(azure): cache the remaining per-request Entra credentials and keep token refresh off the event loop

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-10-19 16:20:20 +00:00
parent 0675c22d56
commit c008f25b55
8 changed files with 268 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

97
litellm/llms/azure/common_utils.py
View file

@ -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)

5
litellm/llms/azure/search/transformation.py
View file

@ -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):

33
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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,

23
litellm/realtime_api/main.py
View file

@ -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,

9
litellm/secret_managers/main.py
View file

@ -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:

21
tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py
View file

@ -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"}

78
tests/unit/llms/azure/test_azure_common_utils.py
View file

@ -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.

50
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -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", '