mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(azure): resolve Azure AD tokens for v1 API version client
When api_version is "v1", "latest", or "preview", litellm creates a standard OpenAI client instead of AzureOpenAI. The OpenAI client does not accept azure_ad_token_provider, so Azure AD authentication was silently dropped, causing auth failures. Now resolves azure_ad_token_provider (by calling it) or azure_ad_token to api_key before constructing the OpenAI client. Fixes https://github.com/BerriAI/litellm/issues/27945 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
c1b7a4b2a3
commit
52fa9da7e5
2 changed files with 87 additions and 1 deletions
|
|
@ -474,8 +474,22 @@ class BaseAzureLLM(BaseOpenAILLM):
|
|||
if self._is_azure_v1_api_version(api_version):
|
||||
# Extract only params that OpenAI client accepts
|
||||
# Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview"
|
||||
v1_api_key = azure_client_params.get("api_key")
|
||||
|
||||
# When using Azure AD auth, resolve the token for the
|
||||
# standard OpenAI client which does not accept
|
||||
# azure_ad_token_provider directly.
|
||||
if not v1_api_key:
|
||||
azure_ad_token_provider = azure_client_params.get(
|
||||
"azure_ad_token_provider"
|
||||
)
|
||||
if azure_ad_token_provider is not None:
|
||||
v1_api_key = azure_ad_token_provider()
|
||||
elif azure_client_params.get("azure_ad_token"):
|
||||
v1_api_key = azure_client_params["azure_ad_token"]
|
||||
|
||||
v1_params = {
|
||||
"api_key": azure_client_params.get("api_key"),
|
||||
"api_key": v1_api_key,
|
||||
"base_url": f"{api_base}/openai/v1/",
|
||||
}
|
||||
if "timeout" in azure_client_params:
|
||||
|
|
|
|||
|
|
@ -1646,6 +1646,78 @@ def test_azure_v1_api_uses_openai_client(api_version):
|
|||
), f"base_url should contain /openai/v1/, got {async_client.base_url}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
|
||||
def test_azure_v1_api_with_azure_ad_token_provider(api_version):
|
||||
"""
|
||||
Test that Azure v1 API versions resolve azure_ad_token_provider to api_key.
|
||||
|
||||
The standard OpenAI client used for v1 API versions does not accept
|
||||
azure_ad_token_provider. When no api_key is provided and an
|
||||
azure_ad_token_provider is set, the provider should be called and the
|
||||
resulting token used as api_key for the OpenAI client.
|
||||
|
||||
See: https://github.com/BerriAI/litellm/issues/27945
|
||||
"""
|
||||
from openai import OpenAI
|
||||
|
||||
base_llm = BaseAzureLLM()
|
||||
api_base = "https://test.openai.azure.com"
|
||||
token_value = "mock-azure-ad-token-from-provider"
|
||||
|
||||
with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
|
||||
mock_init.return_value = {
|
||||
"api_key": None,
|
||||
"azure_endpoint": api_base,
|
||||
"api_version": api_version,
|
||||
"azure_ad_token": None,
|
||||
"azure_ad_token_provider": lambda: token_value,
|
||||
}
|
||||
|
||||
client = base_llm.get_azure_openai_client(
|
||||
api_key=None,
|
||||
api_base=api_base,
|
||||
api_version=api_version,
|
||||
_is_async=False,
|
||||
)
|
||||
|
||||
assert isinstance(client, OpenAI)
|
||||
assert client.api_key == token_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
|
||||
def test_azure_v1_api_with_azure_ad_token(api_version):
|
||||
"""
|
||||
Test that Azure v1 API versions use azure_ad_token as api_key when no
|
||||
api_key or token_provider is set.
|
||||
|
||||
See: https://github.com/BerriAI/litellm/issues/27945
|
||||
"""
|
||||
from openai import OpenAI
|
||||
|
||||
base_llm = BaseAzureLLM()
|
||||
api_base = "https://test.openai.azure.com"
|
||||
token_value = "static-azure-ad-token"
|
||||
|
||||
with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
|
||||
mock_init.return_value = {
|
||||
"api_key": None,
|
||||
"azure_endpoint": api_base,
|
||||
"api_version": api_version,
|
||||
"azure_ad_token": token_value,
|
||||
"azure_ad_token_provider": None,
|
||||
}
|
||||
|
||||
client = base_llm.get_azure_openai_client(
|
||||
api_key=None,
|
||||
api_base=api_base,
|
||||
api_version=api_version,
|
||||
_is_async=False,
|
||||
)
|
||||
|
||||
assert isinstance(client, OpenAI)
|
||||
assert client.api_key == token_value
|
||||
|
||||
|
||||
def test_azure_traditional_api_uses_azure_openai_client():
|
||||
"""
|
||||
Test that traditional Azure API versions still use AzureOpenAI client.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue