feat(azure): Make Azure AD scope configurable (#11621)

This commit is contained in:
Karthick J 2025-06-15 06:13:01 +05:30 committed by GitHub
parent fc57a53a60
commit ac4c29c352
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 34 additions and 5 deletions

View file

@ -558,6 +558,7 @@ model_list:
tenant_id: os.environ/AZURE_TENANT_ID
client_id: os.environ/AZURE_CLIENT_ID
client_secret: os.environ/AZURE_CLIENT_SECRET
azure_scope: os.environ/AZURE_SCOPE # defaults to "https://cognitiveservices.azure.com/.default"
```
Test it
@ -594,6 +595,7 @@ model_list:
client_id: os.environ/AZURE_CLIENT_ID
azure_username: os.environ/AZURE_USERNAME
azure_password: os.environ/AZURE_PASSWORD
azure_scope: os.environ/AZURE_SCOPE # defaults to "https://cognitiveservices.azure.com/.default"
```
Test it

View file

@ -111,6 +111,7 @@ def get_litellm_params(
"client_secret": kwargs.get("client_secret"),
"azure_username": kwargs.get("azure_username"),
"azure_password": kwargs.get("azure_password"),
"azure_scope": kwargs.get("azure_scope"),
"max_retries": max_retries,
"timeout": kwargs.get("timeout"),
"bucket_name": kwargs.get("bucket_name"),

View file

@ -162,6 +162,7 @@ def get_azure_ad_token_from_oidc(
azure_ad_token: str,
azure_client_id: Optional[str],
azure_tenant_id: Optional[str],
scope: str = "https://cognitiveservices.azure.com/.default",
) -> str:
"""
Get Azure AD token from OIDC token
@ -170,6 +171,7 @@ def get_azure_ad_token_from_oidc(
azure_ad_token: str
azure_client_id: Optional[str]
azure_tenant_id: Optional[str]
scope: str
Returns:
`azure_ad_token_access_token` - str
@ -212,7 +214,7 @@ def get_azure_ad_token_from_oidc(
data={
"client_id": azure_client_id,
"grant_type": "client_credentials",
"scope": "https://cognitiveservices.azure.com/.default",
"scope": scope,
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
"client_assertion": oidc_token,
},
@ -335,6 +337,8 @@ class BaseAzureLLM(BaseOpenAILLM):
azure_password = litellm_params.get(
"azure_password", os.getenv("AZURE_PASSWORD")
)
scope = litellm_params.get(
"azure_scope", os.getenv("AZURE_SCOPE", "https://cognitiveservices.azure.com/.default"))
max_retries = litellm_params.get("max_retries")
timeout = litellm_params.get("timeout")
if (
@ -349,6 +353,7 @@ class BaseAzureLLM(BaseOpenAILLM):
tenant_id=tenant_id,
client_id=client_id,
client_secret=client_secret,
scope=scope,
)
if azure_ad_token_provider is None and azure_username and azure_password and client_id:
verbose_logger.debug("Using Azure Username and Password for Azure Auth")
@ -356,6 +361,7 @@ class BaseAzureLLM(BaseOpenAILLM):
azure_username=azure_username,
azure_password=azure_password,
client_id=client_id,
scope=scope,
)
if azure_ad_token is not None and azure_ad_token.startswith("oidc/"):
@ -364,6 +370,7 @@ class BaseAzureLLM(BaseOpenAILLM):
azure_ad_token=azure_ad_token,
azure_client_id=client_id,
azure_tenant_id=tenant_id,
scope=scope,
)
elif (
not api_key
@ -435,6 +442,8 @@ class BaseAzureLLM(BaseOpenAILLM):
## 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"))
scope = litellm_params.get("azure_scope", os.getenv(
"AZURE_SCOPE", "https://cognitiveservices.azure.com/.default"))
if client is None:
if not api_base.endswith("/"):
api_base += "/"
@ -455,6 +464,7 @@ class BaseAzureLLM(BaseOpenAILLM):
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

View file

@ -1261,6 +1261,7 @@ def completion( # type: ignore # noqa: PLR0915
client_secret=kwargs.get("client_secret"),
azure_username=kwargs.get("azure_username"),
azure_password=kwargs.get("azure_password"),
azure_scope=kwargs.get("azure_scope"),
max_retries=max_retries,
timeout=timeout,
)

View file

@ -2153,6 +2153,7 @@ all_litellm_params = [
"client_id",
"azure_username",
"azure_password",
"azure_scope",
"client_secret",
"user_continue_message",
"configurable_clientside_auth_params",

View file

@ -82,7 +82,8 @@ def test_initialize_with_tenant_credentials_env_var(setup_mocks, monkeypatch):
monkeypatch.setenv("AZURE_TENANT_ID", "test-tenant-id")
monkeypatch.setenv("AZURE_CLIENT_ID", "test-client-id")
monkeypatch.setenv("AZURE_CLIENT_SECRET", "test-client-secret")
monkeypatch.setenv("AZURE_SCOPE", "test-azure-scope")
result = BaseAzureLLM().initialize_azure_sdk_client(
litellm_params={},
api_key=None,
@ -97,6 +98,7 @@ def test_initialize_with_tenant_credentials_env_var(setup_mocks, monkeypatch):
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
scope="test-azure-scope"
)
# Verify expected result
@ -112,6 +114,7 @@ def test_initialize_with_tenant_credentials(setup_mocks):
"tenant_id": "test-tenant-id",
"client_id": "test-client-id",
"client_secret": "test-client-secret",
"azure_scope": "test-azure-scope",
},
api_key=None,
api_base="https://test.openai.azure.com",
@ -125,6 +128,7 @@ def test_initialize_with_tenant_credentials(setup_mocks):
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
scope="test-azure-scope",
)
# Verify expected result
@ -139,6 +143,7 @@ def test_initialize_with_username_password(monkeypatch, setup_mocks):
monkeypatch.delenv("AZURE_CLIENT_SECRET", raising=False)
monkeypatch.delenv("AZURE_USERNAME", raising=False)
monkeypatch.delenv("AZURE_PASSWORD", raising=False)
monkeypatch.delenv("AZURE_SCOPE", raising=False)
# Test with azure_username, azure_password, and client_id provided
result = BaseAzureLLM().initialize_azure_sdk_client(
@ -146,6 +151,7 @@ def test_initialize_with_username_password(monkeypatch, setup_mocks):
"azure_username": "test-username",
"azure_password": "test-password",
"client_id": "test-client-id",
"azure_scope": "test-azure-scope"
},
api_key=None,
api_base="https://test.openai.azure.com",
@ -167,6 +173,7 @@ def test_initialize_with_username_password(monkeypatch, setup_mocks):
azure_username="test-username",
azure_password="test-password",
client_id="test-client-id",
scope="test-azure-scope"
)
# Verify expected result
@ -176,6 +183,8 @@ def test_initialize_with_username_password(monkeypatch, 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)
monkeypatch.delenv("AZURE_SCOPE", 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"},
@ -186,9 +195,9 @@ def test_initialize_with_oidc_token(setup_mocks, monkeypatch):
is_async=False,
)
# Verify that get_azure_ad_token_from_oidc was called
setup_mocks["oidc_token"].assert_called_once_with(
azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None
azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None,
scope="https://cognitiveservices.azure.com/.default"
)
# Verify expected result
@ -202,6 +211,7 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks):
"azure_ad_token": "oidc/test-token",
"client_id": "test-client-id",
"tenant_id": "test-tenant-id",
"azure_scope": "test-azure-scope",
},
api_key=None,
api_base="https://test.openai.azure.com",
@ -215,6 +225,7 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks):
azure_ad_token="oidc/test-token",
azure_client_id="test-client-id",
azure_tenant_id="test-tenant-id",
scope="test-azure-scope"
)
# Verify expected result
@ -243,6 +254,7 @@ def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch):
azure_ad_token="oidc/test-token",
azure_client_id="env-client-id",
azure_tenant_id="env-tenant-id",
scope="https://cognitiveservices.azure.com/.default"
)
# Verify expected result
@ -253,6 +265,7 @@ 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)
monkeypatch.delenv("AZURE_SCOPE", raising=False)
# Test with azure_ad_token that starts with "oidc/" but no credentials anywhere
result = BaseAzureLLM().initialize_azure_sdk_client(
@ -268,7 +281,8 @@ def test_initialize_with_oidc_token_no_credentials(setup_mocks, monkeypatch):
# 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
azure_ad_token="oidc/test-token", azure_client_id=None, azure_tenant_id=None,
scope="https://cognitiveservices.azure.com/.default"
)
# Verify expected result