scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

[Fixes] Azure OpenAI OIDC - allow using litellm defined params for OIDC Auth (#10394)

Browse source 
* fixes for azure OIDC ad token auth

* fixes for test_azure_common_utils.py

* get_azure_ad_token_from_oidc
This commit is contained in:
Ishaan Jaff 2025-05-01 14:46:56 -07:00 • committed by GitHub
parent 03718b8b76
commit eebb725416
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 124 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/llms/azure/azure.py
View file

@ -125,22 +125,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
def __init__(self) -> None:
super().__init__()
def validate_environment(self, api_key, azure_ad_token, azure_ad_token_provider):
headers = {
"content-type": "application/json",
}
if api_key is not None:
headers["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_ad_token)
headers["Authorization"] = f"Bearer {azure_ad_token}"
elif azure_ad_token_provider is not None:
azure_ad_token = azure_ad_token_provider()
headers["Authorization"] = f"Bearer {azure_ad_token}"
return headers
def make_sync_azure_openai_chat_completion_request(
self,
azure_client: AzureOpenAI,
@ -242,6 +226,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
azure_ad_token_provider=azure_ad_token_provider,
acompletion=acompletion,
client=client,
litellm_params=litellm_params,
)
data = {"model": None, "messages": messages, **optional_params}

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

@ -158,13 +158,27 @@ def get_azure_ad_token_from_username_password(
return token_provider
def get_azure_ad_token_from_oidc(azure_ad_token: str):
azure_client_id = os.getenv("AZURE_CLIENT_ID", None)
azure_tenant_id = os.getenv("AZURE_TENANT_ID", None)
def get_azure_ad_token_from_oidc(
azure_ad_token: str,
azure_client_id: Optional[str],
azure_tenant_id: Optional[str],
) -> str:
"""
Get Azure AD token from OIDC token
Args:
azure_ad_token: str
azure_client_id: Optional[str]
azure_tenant_id: Optional[str]
Returns:
`azure_ad_token_access_token` - str
"""
azure_authority_host = os.getenv(
"AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com"
)
azure_client_id = azure_client_id or os.getenv("AZURE_CLIENT_ID")
azure_tenant_id = azure_tenant_id or os.getenv("AZURE_TENANT_ID")
if azure_client_id is None or azure_tenant_id is None:
raise AzureOpenAIError(
status_code=422,
@ -341,7 +355,11 @@ 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)
azure_ad_token = get_azure_ad_token_from_oidc(
azure_ad_token=azure_ad_token,
azure_client_id=client_id,
azure_tenant_id=tenant_id,
)
elif (
not api_key
and azure_ad_token_provider is None
@ -402,6 +420,7 @@ class BaseAzureLLM(BaseOpenAILLM):
api_version: str,
max_retries: int,
timeout: Union[float, httpx.Timeout],
litellm_params: dict,
api_key: Optional[str],
azure_ad_token: Optional[str],
azure_ad_token_provider: Optional[Callable[[], str]],
@ -409,6 +428,8 @@ class BaseAzureLLM(BaseOpenAILLM):
client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None,
) -> Union[AzureOpenAI, AsyncAzureOpenAI]:
## build base url - assume api base includes resource name
tenant_id = litellm_params.get("tenant_id", os.getenv("AZURE_TENANT_ID"))
client_id = litellm_params.get("client_id", os.getenv("AZURE_CLIENT_ID"))
if client is None:
if not api_base.endswith("/"):
api_base += "/"
@ -425,7 +446,11 @@ 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_ad_token)
azure_ad_token = get_azure_ad_token_from_oidc(
azure_ad_token=azure_ad_token,
azure_client_id=client_id,
azure_tenant_id=tenant_id,
)
azure_client_params["azure_ad_token"] = azure_ad_token
if azure_ad_token_provider is not None:

1
litellm/llms/azure/completion/handler.py
View file

@ -72,6 +72,7 @@ class AzureTextCompletion(BaseAzureLLM):
azure_ad_token=azure_ad_token,
azure_ad_token_provider=azure_ad_token_provider,
acompletion=acompletion,
litellm_params=litellm_params,
)
data = {"model": None, "prompt": prompt, **optional_params}

88
tests/litellm/llms/azure/test_azure_common_utils.py
View file

@ -159,7 +159,9 @@ def test_initialize_with_username_password(setup_mocks):
assert "azure_ad_token_provider" in result
def test_initialize_with_oidc_token(setup_mocks):
def test_initialize_with_oidc_token(setup_mocks, monkeypatch):
monkeypatch.delenv("AZURE_CLIENT_ID", raising=False)
monkeypatch.delenv("AZURE_TENANT_ID", raising=False)
# Test with azure_ad_token that starts with "oidc/"
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={"azure_ad_token": "oidc/test-token"},
@ -171,7 +173,89 @@ def test_initialize_with_oidc_token(setup_mocks):
)
# Verify that get_azure_ad_token_from_oidc was called
setup_mocks["oidc_token"].assert_called_once_with("oidc/test-token")
setup_mocks["oidc_token"].assert_called_once_with(
azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None
)
# 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
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={
"azure_ad_token": "oidc/test-token",
"client_id": "test-client-id",
"tenant_id": "test-tenant-id",
},
api_key=None,
api_base="https://test.openai.azure.com",
model_name="gpt-4",
api_version=None,
is_async=False,
)
# Verify that get_azure_ad_token_from_oidc was called with the correct parameters
setup_mocks["oidc_token"].assert_called_once_with(
azure_ad_token="oidc/test-token",
azure_client_id="test-client-id",
azure_tenant_id="test-tenant-id",
)
# 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
monkeypatch.setenv("AZURE_CLIENT_ID", "env-client-id")
monkeypatch.setenv("AZURE_TENANT_ID", "env-tenant-id")
# Test with azure_ad_token that starts with "oidc/" but no explicit client/tenant IDs
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={
"azure_ad_token": "oidc/test-token",
},
api_key=None,
api_base="https://test.openai.azure.com",
model_name="gpt-4",
api_version=None,
is_async=False,
)
# Verify that get_azure_ad_token_from_oidc was called with environment variables
setup_mocks["oidc_token"].assert_called_once_with(
azure_ad_token="oidc/test-token",
azure_client_id="env-client-id",
azure_tenant_id="env-tenant-id",
)
# Verify expected result
assert result["azure_ad_token"] == "mock-oidc-token"
def test_initialize_with_oidc_token_no_credentials(setup_mocks, monkeypatch):
# Clear environment variables
monkeypatch.delenv("AZURE_CLIENT_ID", raising=False)
monkeypatch.delenv("AZURE_TENANT_ID", raising=False)
# Test with azure_ad_token that starts with "oidc/" but no credentials anywhere
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={
"azure_ad_token": "oidc/test-token",
},
api_key=None,
api_base="https://test.openai.azure.com",
model_name="gpt-4",
api_version=None,
is_async=False,
)
# Verify that get_azure_ad_token_from_oidc was called with None values
setup_mocks["oidc_token"].assert_called_once_with(
azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None
)
# Verify expected result
assert result["azure_ad_token"] == "mock-oidc-token"

6
tests/litellm_utils_tests/test_secret_manager.py
View file

@ -140,7 +140,11 @@ def test_oidc_circleci_with_azure():
# TODO: Switch to our own Azure account, currently using ai.moda's account
os.environ["AZURE_TENANT_ID"] = "17c0a27a-1246-4aa1-a3b6-d294e80e783c"
os.environ["AZURE_CLIENT_ID"] = "4faf5422-b2bd-45e8-a6d7-46543a38acd0"
azure_ad_token = get_azure_ad_token_from_oidc("oidc/circleci/")
azure_ad_token = get_azure_ad_token_from_oidc(
azure_ad_token="oidc/circleci/",
azure_client_id=None,
azure_tenant_id=None,
)
print(f"secret_val: {redact_oidc_signature(azure_ad_token)}")