From ac4c29c352eff187158fa04896e823ab491a96df Mon Sep 17 00:00:00 2001 From: Karthick J Date: Sun, 15 Jun 2025 06:13:01 +0530 Subject: [PATCH] feat(azure): Make Azure AD scope configurable (#11621) --- docs/my-website/docs/providers/azure/azure.md | 2 ++ .../litellm_core_utils/get_litellm_params.py | 1 + litellm/llms/azure/common_utils.py | 12 +++++++++- litellm/main.py | 1 + litellm/types/utils.py | 1 + .../llms/azure/test_azure_common_utils.py | 22 +++++++++++++++---- 6 files changed, 34 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/providers/azure/azure.md b/docs/my-website/docs/providers/azure/azure.md index d0b0371986..065654df6f 100644 --- a/docs/my-website/docs/providers/azure/azure.md +++ b/docs/my-website/docs/providers/azure/azure.md @@ -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 diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 19c8ec8d80..c354dea024 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -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"), diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 3238b8e862..2d855c476d 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -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 diff --git a/litellm/main.py b/litellm/main.py index fea52bb97b..59a7764d31 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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, ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index affe06bab1..fd8ecd9e14 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2153,6 +2153,7 @@ all_litellm_params = [ "client_id", "azure_username", "azure_password", + "azure_scope", "client_secret", "user_continue_message", "configurable_clientside_auth_params", diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index d7e9e983f9..9012fe6c72 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -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