scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 150

Merge 6813c875e0 into bc18e10f6b

Browse source
This commit is contained in:
Cedric Specht 2026-10-23 17:17:40 +00:00 • committed by GitHub
commit ae08cb4ef9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 226 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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