diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 23c5aa35f4b..901d2c965db 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -5,7 +5,8 @@ import os import re from collections.abc import Awaitable, Callable, Mapping from contextlib import closing -from functools import lru_cache +from functools import lru_cache, partial +from time import monotonic from types import MappingProxyType from typing import Final, Literal, NamedTuple, cast @@ -30,6 +31,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.utils import add_path_to_api_base azure_ad_cache: Final = DualCache() +_AZURE_AD_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60 _LAST_OPENAI_PATH_SEGMENT: Final = re.compile(r"^(.*/openai)(?=/|$)") @@ -283,6 +285,7 @@ def get_azure_ad_token_from_oidc( "azure_tenant_id": azure_tenant_id, "azure_authority_host": azure_authority_host, "oidc_token": oidc_token, + "scope": scope, } ) @@ -290,6 +293,7 @@ def get_azure_ad_token_from_oidc( if azure_ad_token_access_token is not None: return azure_ad_token_access_token + request_started_at: Final = monotonic() with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: req_token: Final = client.post( f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", @@ -318,15 +322,34 @@ def get_azure_ad_token_from_oidc( if azure_ad_token_expires_in is None: raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") - azure_ad_cache.set_cache( - key=azure_ad_token_cache_key, - value=azure_ad_token_access_token, - ttl=azure_ad_token_expires_in, + cache_ttl: Final = ( + azure_ad_token_expires_in - _AZURE_AD_TOKEN_EXPIRY_MARGIN_SECONDS - (monotonic() - request_started_at) ) + if cache_ttl > 0: + azure_ad_cache.set_cache( + key=azure_ad_token_cache_key, + value=azure_ad_token_access_token, + ttl=cache_ttl, + ) return azure_ad_token_access_token +def get_azure_ad_token_provider_from_oidc( + azure_ad_token: str, + azure_client_id: str | None, + azure_tenant_id: str | None, + scope: str | None, +) -> Callable[[], str]: + return partial( + get_azure_ad_token_from_oidc, + azure_ad_token=azure_ad_token, + azure_client_id=azure_client_id, + azure_tenant_id=azure_tenant_id, + scope=scope, + ) + + def select_azure_base_url_or_endpoint(azure_client_params: dict): azure_endpoint: Final = azure_client_params.get("azure_endpoint", None) if azure_endpoint is not None: @@ -704,15 +727,22 @@ class BaseAzureLLM(BaseOpenAILLM): scope=scope, ) - 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( + oidc_token_provider: Final = ( + get_azure_ad_token_provider_from_oidc( azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope, ) - elif not api_key and azure_ad_token_provider is None and litellm.enable_azure_ad_token_refresh is True: + if azure_ad_token is not None and azure_ad_token.startswith("oidc/") + else None + ) + if ( + oidc_token_provider is None + and 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" ) @@ -734,8 +764,8 @@ class BaseAzureLLM(BaseOpenAILLM): "api_key": api_key, "azure_endpoint": api_base, "api_version": api_version, - "azure_ad_token": azure_ad_token, - "azure_ad_token_provider": azure_ad_token_provider, + "azure_ad_token": None if oidc_token_provider is not None else azure_ad_token, + "azure_ad_token_provider": azure_ad_token_provider or oidc_token_provider, } # init http client + SSL Verification settings if is_async is True: @@ -797,14 +827,14 @@ class BaseAzureLLM(BaseOpenAILLM): azure_client_params["api_key"] = api_key elif azure_ad_token is not None: if azure_ad_token.startswith("oidc/"): - azure_ad_token = get_azure_ad_token_from_oidc( + azure_client_params["azure_ad_token_provider"] = get_azure_ad_token_provider_from_oidc( azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope, ) - - azure_client_params["azure_ad_token"] = azure_ad_token + else: + azure_client_params["azure_ad_token"] = azure_ad_token if azure_ad_token_provider is not None: azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 1dc7ff74323..67fcb591003 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +import respx import litellm from litellm.llms.azure.common_utils import ( @@ -224,6 +225,8 @@ def test_initialize_with_oidc_token(setup_mocks, monkeypatch): is_async=False, ) + assert result["azure_ad_token"] is None + 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, @@ -231,9 +234,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): # Test with azure_ad_token that starts with "oidc/" and explicit client/tenant IDs @@ -251,7 +251,8 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks): is_async=False, ) - # Verify that get_azure_ad_token_from_oidc was called with the correct parameters + assert result["azure_ad_token"] is None + 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", @@ -259,9 +260,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): # Set environment variables @@ -280,7 +278,8 @@ def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch): is_async=False, ) - # Verify that get_azure_ad_token_from_oidc was called with environment variables + assert result["azure_ad_token"] is None + 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", @@ -288,9 +287,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): # Clear environment variables @@ -2360,3 +2356,178 @@ def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(m closer.reap() assert wrapper.is_closed() is True + + +class OidcTokenService: + def __init__(self) -> None: + self.seconds = 0.0 + self.expires_in = 3600 + self.exchange_delay = 0.0 + self.scopes: tuple[str, ...] = () + self.authorization_headers: tuple[str, ...] = () + + def now(self) -> float: + return self.seconds + + def exchange(self, request: httpx.Request) -> httpx.Response: + scope: Final = parse_qs(request.content.decode())["scope"][0] + self.scopes += (scope,) + self.seconds += self.exchange_delay + return httpx.Response(200, json={"access_token": f"token-{len(self.scopes)}", "expires_in": self.expires_in}) + + def completion(self, request: httpx.Request) -> httpx.Response: + self.authorization_headers += (request.headers["Authorization"],) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "test-deployment", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + }, + ) + + +@pytest.fixture +def oidc_service(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> OidcTokenService: + from litellm.caching.caching import DualCache + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.llms.azure import common_utils + from litellm.secret_managers import main as secret_manager + + service: Final = OidcTokenService() + monkeypatch.setenv("AZURE_CLIENT_ID", "test-client") + monkeypatch.setenv("AZURE_TENANT_ID", "test-tenant") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + monkeypatch.delenv("AZURE_CLIENT_SECRET", raising=False) + monkeypatch.setattr(common_utils, "azure_ad_cache", DualCache(in_memory_cache=InMemoryCache(clock=service.now))) + monkeypatch.setattr(common_utils, "monotonic", service.now, raising=False) + monkeypatch.setattr(secret_manager, "oidc_cache", DualCache(in_memory_cache=InMemoryCache(clock=service.now))) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", InMemoryCache(clock=service.now)) + monkeypatch.setattr(litellm, "client_session", None) + monkeypatch.setattr(litellm, "aclient_session", None) + respx_mock.get( + "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity" + ).respond(200, text="synthetic-google-assertion") + respx_mock.post("https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token").mock( + side_effect=service.exchange + ) + respx_mock.route(method="POST", host="test.openai.azure.com").mock(side_effect=service.completion) + return service + + +def test_oidc_tokens_are_isolated_by_scope(oidc_service: OidcTokenService) -> None: + reference: Final = "oidc/google/api://AzureADTokenExchange" + cognitive_token: Final = get_azure_ad_token_from_oidc(reference) + ai_token: Final = get_azure_ad_token_from_oidc(reference, scope="https://ai.azure.com/.default") + + assert cognitive_token == "token-1" + assert ai_token == "token-2" + assert get_azure_ad_token_from_oidc(reference) == cognitive_token + assert oidc_service.scopes == ("https://cognitiveservices.azure.com/.default", "https://ai.azure.com/.default") + + +def test_oidc_refresh_margin_includes_exchange_time(oidc_service: OidcTokenService) -> None: + reference: Final = "oidc/google/api://AzureADTokenExchange" + oidc_service.exchange_delay = 20 + assert get_azure_ad_token_from_oidc(reference) == "token-1" + oidc_service.seconds = 3539 + assert get_azure_ad_token_from_oidc(reference) == "token-1" + oidc_service.seconds = 3541 + assert get_azure_ad_token_from_oidc(reference) == "token-2" + + +@pytest.mark.parametrize("expires_in", [1, 60]) +def test_oidc_short_lived_tokens_are_not_cached(oidc_service: OidcTokenService, expires_in: int) -> None: + oidc_service.expires_in = expires_in + reference: Final = "oidc/google/api://AzureADTokenExchange" + assert get_azure_ad_token_from_oidc(reference) == "token-1" + assert get_azure_ad_token_from_oidc(reference) == "token-2" + + +@pytest.mark.parametrize("api_version", ["2024-10-21", "v1"]) +def test_oidc_cached_sdk_client_refreshes_authorization(oidc_service: OidcTokenService, api_version: str) -> None: + from openai import AzureOpenAI, OpenAI + + client: Final = BaseAzureLLM().get_azure_openai_client( + api_key=None, + api_base="https://test.openai.azure.com", + api_version=api_version, + model="test-deployment", + litellm_params={"azure_ad_token": "oidc/google/api://AzureADTokenExchange"}, + ) + assert isinstance(client, (AzureOpenAI, OpenAI)) + with client: + client.chat.completions.create(model="test-deployment", messages=[{"role": "user", "content": "Hi"}]) + oidc_service.seconds = 3590 + client.chat.completions.create(model="test-deployment", messages=[{"role": "user", "content": "Hi"}]) + assert oidc_service.authorization_headers == ("Bearer token-1", "Bearer token-2") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("api_version", ["2024-10-21", "v1"]) +async def test_oidc_cached_async_sdk_client_refreshes_authorization( + oidc_service: OidcTokenService, api_version: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from openai import AsyncAzureOpenAI, AsyncOpenAI + + http_client: Final = httpx.AsyncClient(transport=httpx.MockTransport(oidc_service.completion)) + monkeypatch.setattr(litellm, "aclient_session", http_client) + client: Final = BaseAzureLLM().get_azure_openai_client( + api_key=None, + api_base="https://test.openai.azure.com", + api_version=api_version, + model="test-deployment", + litellm_params={"azure_ad_token": "oidc/google/api://AzureADTokenExchange"}, + _is_async=True, + ) + assert isinstance(client, (AsyncAzureOpenAI, AsyncOpenAI)) + async with client: + await client.chat.completions.create(model="test-deployment", messages=[{"role": "user", "content": "Hi"}]) + oidc_service.seconds = 3590 + await client.chat.completions.create(model="test-deployment", messages=[{"role": "user", "content": "Hi"}]) + assert oidc_service.authorization_headers == ("Bearer token-1", "Bearer token-2") + + +def test_oidc_shared_headers_refresh_and_keep_scopes_separate(oidc_service: OidcTokenService) -> None: + cognitive_params: Final = GenericLiteLLMParams(azure_ad_token="oidc/google/api://AzureADTokenExchange") + ai_params: Final = GenericLiteLLMParams( + azure_ad_token="oidc/google/api://AzureADTokenExchange", azure_scope="https://ai.azure.com/.default" + ) + assert BaseAzureLLM._base_validate_azure_environment({}, cognitive_params)["Authorization"] == "Bearer token-1" + assert BaseAzureLLM._base_validate_azure_environment({}, ai_params)["Authorization"] == "Bearer token-2" + oidc_service.seconds = 3590 + assert BaseAzureLLM._base_validate_azure_environment({}, cognitive_params)["Authorization"] == "Bearer token-3" + + +@pytest.mark.parametrize( + ("token", "expected_headers"), + [ + ("oidc/google/api://AzureADTokenExchange", ("Bearer token-1", "Bearer token-2")), + ("static-token", ("Bearer static-token", "Bearer static-token")), + ], +) +def test_oidc_cloudflare_client_refreshes_without_changing_static_tokens( + oidc_service: OidcTokenService, token: str, expected_headers: tuple[str, str] +) -> None: + from openai import AzureOpenAI + + client: Final = BaseAzureLLM()._init_azure_client_for_cloudflare_ai_gateway( + api_base="https://test.openai.azure.com", + model="test-deployment", + api_version="2024-10-21", + max_retries=0, + timeout=30, + litellm_params={}, + api_key=None, + azure_ad_token=token, + azure_ad_token_provider=None, + acompletion=False, + ) + assert isinstance(client, AzureOpenAI) + with client: + client.chat.completions.create(model="test-deployment", messages=[{"role": "user", "content": "Hi"}]) + oidc_service.seconds = 3590 + client.chat.completions.create(model="test-deployment", messages=[{"role": "user", "content": "Hi"}]) + assert oidc_service.authorization_headers == expected_headers