diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index cce2265d280..c035a7401bd 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -4,7 +4,7 @@ import json import os import re from collections.abc import Awaitable, Callable, Mapping -from functools import lru_cache +from functools import lru_cache, partial from types import MappingProxyType from typing import Final, Literal, NamedTuple, cast @@ -703,12 +703,14 @@ class BaseAzureLLM(BaseOpenAILLM): if azure_ad_token is not None and azure_ad_token.startswith("oidc/"): verbose_logger.debug("Using Azure OIDC Token for Azure Auth") - azure_ad_token = get_azure_ad_token_from_oidc( + azure_ad_token_provider = partial( + get_azure_ad_token_from_oidc, azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope, ) + azure_ad_token = None elif not api_key and azure_ad_token_provider is None and litellm.enable_azure_ad_token_refresh is True: verbose_logger.debug( "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth" diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index c60e3e8a9d8..3f2b63c74a3 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -223,6 +223,7 @@ def test_initialize_with_oidc_token(setup_mocks, monkeypatch): is_async=False, ) + assert result["azure_ad_token_provider"]() == "mock-oidc-token" setup_mocks["oidc_token"].assert_called_once_with( azure_ad_token="oidc/test-token", azure_client_id=None, @@ -230,8 +231,6 @@ def test_initialize_with_oidc_token(setup_mocks, monkeypatch): scope="https://cognitiveservices.azure.com/.default", ) - # Verify expected result - assert result["azure_ad_token"] == "mock-oidc-token" def test_initialize_with_oidc_token_and_client_params(setup_mocks): @@ -251,6 +250,7 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks): ) # Verify that get_azure_ad_token_from_oidc was called with the correct parameters + assert result["azure_ad_token_provider"]() == "mock-oidc-token" setup_mocks["oidc_token"].assert_called_once_with( azure_ad_token="oidc/test-token", azure_client_id="test-client-id", @@ -258,8 +258,6 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks): scope="test-azure-scope", ) - # Verify expected result - assert result["azure_ad_token"] == "mock-oidc-token" def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch): @@ -280,6 +278,7 @@ def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch): ) # Verify that get_azure_ad_token_from_oidc was called with environment variables + assert result["azure_ad_token_provider"]() == "mock-oidc-token" setup_mocks["oidc_token"].assert_called_once_with( azure_ad_token="oidc/test-token", azure_client_id="env-client-id", @@ -287,8 +286,6 @@ def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch): scope="https://cognitiveservices.azure.com/.default", ) - # Verify expected result - assert result["azure_ad_token"] == "mock-oidc-token" def test_initialize_with_ad_token_provider(setup_mocks, monkeypatch): @@ -2217,6 +2214,69 @@ 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_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_request(monkeypatch, api_version): + import threading + + import httpx + + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + exchange_threads: Final[list[int]] = [] + + def entra_token_endpoint(request: httpx.Request) -> httpx.Response: + exchange_threads.append(threading.get_ident()) + assert b"client_assertion=federated-oidc-token" in request.content + return httpx.Response(200, json={"access_token": f"exchanged-{api_version}", "expires_in": 3600}) + + 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"]}, + } + ], + }, + ) + + monkeypatch.setenv("FEDERATED_OIDC_TOKEN", "federated-oidc-token") + monkeypatch.setattr( + litellm, + "module_level_client", + HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(entra_token_endpoint))), + ) + monkeypatch.setattr( + litellm, "aclient_session", httpx.AsyncClient(transport=httpx.MockTransport(echo_authorization)) + ) + loop_thread: Final = threading.get_ident() + client: Final = BaseAzureLLM().get_azure_openai_client( + api_key=None, + api_base=f"https://oidc-exchange-{api_version}.openai.azure.com", + api_version=api_version, + litellm_params={ + "azure_ad_token": "oidc/env/FEDERATED_OIDC_TOKEN", + "client_id": f"oidc-client-{api_version}", + "tenant_id": "oidc-tenant", + }, + _is_async=True, + ) + + response: Final = await client.chat.completions.create(model="gpt", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == f"Bearer exchanged-{api_version}" + assert len(exchange_threads) == 1 + assert loop_thread not in exchange_threads + + @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):