mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
feat(azure): Make Azure AD scope configurable (#11621)
This commit is contained in:
parent
fc57a53a60
commit
ac4c29c352
6 changed files with 34 additions and 5 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2153,6 +2153,7 @@ all_litellm_params = [
|
|||
"client_id",
|
||||
"azure_username",
|
||||
"azure_password",
|
||||
"azure_scope",
|
||||
"client_secret",
|
||||
"user_continue_message",
|
||||
"configurable_clientside_auth_params",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue