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 267

fix(azure): await async azure_ad_token_provider results in the threaded wrapper

Browse source
This commit is contained in:
mateo-berri 2026-10-19 09:52:25 -07:00
parent 4848674037
commit 7103bfd003
3 changed files with 63 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -46,7 +46,6 @@ class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object):
_ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams)
_AZURE_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter(Callable[[], str])
class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object):
@ -136,10 +135,11 @@ def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]:
def _async_azure_ad_token_provider(
sync_provider: Callable[[], str],
provider: Callable[[], str | Awaitable[str]],
) -> Callable[[], Awaitable[str]]:
async def async_provider() -> str:
return await asyncio.to_thread(sync_provider)
token: Final = await asyncio.to_thread(provider)
return token if isinstance(token, str) else await token
return async_provider
@ -607,7 +607,7 @@ class BaseAzureLLM(BaseOpenAILLM):
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))
_async_azure_ad_token_provider(v1_token_provider)
if _is_async is True and callable(v1_token_provider)
else v1_token_provider
)
@ -639,11 +639,9 @@ class BaseAzureLLM(BaseOpenAILLM):
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)
),
"azure_ad_token_provider": _async_azure_ad_token_provider(async_token_provider_candidate),
}
if callable(async_token_provider_candidate)
if async_token_provider_candidate is not None
else async_azure_client_params
)
openai_client = AsyncAzureOpenAI(**async_azure_client_params_with_threaded_provider)

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

@ -3148,7 +3148,9 @@ class BaseLLMHTTPHandler:
headers: Final = await asyncio.to_thread(
responses_api_provider_config.validate_environment,
headers=extra_headers or {}, model="None", litellm_params=litellm_params
headers=extra_headers or {},
model="None",
litellm_params=litellm_params,
)
if extra_headers:
@ -3403,7 +3405,9 @@ class BaseLLMHTTPHandler:
headers: Final = await asyncio.to_thread(
responses_api_provider_config.validate_environment,
headers=extra_headers or {}, model="None", litellm_params=litellm_params
headers=extra_headers or {},
model="None",
litellm_params=litellm_params,
)
if extra_headers:
@ -3568,7 +3572,9 @@ class BaseLLMHTTPHandler:
headers: Final = await asyncio.to_thread(
responses_api_provider_config.validate_environment,
headers=extra_headers or {}, model="None", litellm_params=litellm_params
headers=extra_headers or {},
model="None",
litellm_params=litellm_params,
)
if extra_headers:
@ -4626,7 +4632,9 @@ class BaseLLMHTTPHandler:
headers: Final = await asyncio.to_thread(
responses_api_provider_config.validate_environment,
headers=extra_headers or {}, model="None", litellm_params=litellm_params
headers=extra_headers or {},
model="None",
litellm_params=litellm_params,
)
if extra_headers:
@ -4802,7 +4810,9 @@ class BaseLLMHTTPHandler:
validated_headers: Final = await asyncio.to_thread(
responses_api_provider_config.validate_environment,
headers=extra_headers or {}, model=model, litellm_params=litellm_params
headers=extra_headers or {},
model=model,
litellm_params=litellm_params,
)
if extra_headers:

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

@ -2217,6 +2217,48 @@ async def test_non_v1_azure_token_refresh_does_not_block_the_event_loop():
await client.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("api_version", ["2024-02-01", "v1"])
async def test_async_azure_ad_token_provider_reaches_azure_as_bearer_token(monkeypatch, api_version):
import httpx
def echo_authorization(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 0,
"model": "gpt",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {"role": "assistant", "content": request.headers["authorization"]},
}
],
},
)
async def async_token_provider() -> str:
return "async-token"
monkeypatch.setattr(
litellm, "aclient_session", httpx.AsyncClient(transport=httpx.MockTransport(echo_authorization))
)
client: Final = BaseAzureLLM().get_azure_openai_client(
api_key=None,
api_base=f"https://async-token-provider-{api_version}.openai.azure.com",
api_version=api_version,
litellm_params={"azure_ad_token_provider": async_token_provider},
_is_async=True,
)
response: Final = await client.chat.completions.create(model="gpt", messages=[{"role": "user", "content": "hi"}])
assert response.choices[0].message.content == "Bearer async-token"
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.