feat(secret_managers): support Azure Workload Identity for Key Vault auth

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm/proxy/proxy_server.py

@@ -126,6 +126,7 @@ from litellm.utils import (
 
 if TYPE_CHECKING:
     from aiohttp import ClientSession
+    from azure.core.credentials import TokenCredential
     from opentelemetry.trace import Span as _Span
 
     from litellm.integrations.opentelemetry import OpenTelemetry
@@ -2040,12 +2041,33 @@ def _resolve_pydantic_type(typ) -> List:
     return typs
 
 
+def _get_azure_key_vault_credential() -> "TokenCredential":
+    """
+    Build the credential used to authenticate to Azure Key Vault.
+
+    When ``AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY`` is truthy, use
+    ``WorkloadIdentityCredential`` (federated token auth for AKS workload identity,
+    no client secret needed); it reads ``AZURE_TENANT_ID``, ``AZURE_CLIENT_ID`` and
+    ``AZURE_FEDERATED_TOKEN_FILE`` from the environment. Otherwise fall back to
+    ``DefaultAzureCredential``.
+    """
+    from litellm.secret_managers.main import str_to_bool
+
+    if str_to_bool(os.getenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY")) is True:
+        from azure.identity import WorkloadIdentityCredential
+
+        return WorkloadIdentityCredential()
+
+    from azure.identity import DefaultAzureCredential
+
+    return DefaultAzureCredential()
+
+
 def load_from_azure_key_vault(use_azure_key_vault: bool = False):
     if use_azure_key_vault is False:
         return
 
     try:
-        from azure.identity import DefaultAzureCredential
         from azure.keyvault.secrets import SecretClient
 
         # Set your Azure Key Vault URI
@@ -2054,7 +2076,7 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False):
         if KVUri is None:
             raise Exception("Error when loading keys from Azure Key Vault: AZURE_KEY_VAULT_URI is not set.")
 
-        credential = DefaultAzureCredential()
+        credential = _get_azure_key_vault_credential()
 
         # Create the SecretClient using the credential
         client = SecretClient(vault_url=KVUri, credential=credential)

tests/test_litellm/proxy/proxy_server/test_lifecycle.py

@@ -417,6 +417,60 @@ def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch)
     assert result is None
 
 
+def _install_fake_azure_identity(monkeypatch):
+    """Install a fake ``azure.identity`` so credential selection can be exercised
+    without the real Azure SDK installed. Returns the two sentinel classes."""
+    import sys
+    import types
+
+    class FakeDefaultAzureCredential:
+        pass
+
+    class FakeWorkloadIdentityCredential:
+        pass
+
+    azure_module = types.ModuleType("azure")
+    identity_module = types.ModuleType("azure.identity")
+    identity_module.DefaultAzureCredential = FakeDefaultAzureCredential
+    identity_module.WorkloadIdentityCredential = FakeWorkloadIdentityCredential
+    azure_module.identity = identity_module
+
+    monkeypatch.setitem(sys.modules, "azure", azure_module)
+    monkeypatch.setitem(sys.modules, "azure.identity", identity_module)
+    return FakeDefaultAzureCredential, FakeWorkloadIdentityCredential
+
+
+def test_get_azure_key_vault_credential_defaults_to_default_credential(monkeypatch):
+    default_cls, workload_cls = _install_fake_azure_identity(monkeypatch)
+    monkeypatch.delenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY", raising=False)
+
+    credential = ps._get_azure_key_vault_credential()
+
+    assert isinstance(credential, default_cls)
+    assert not isinstance(credential, workload_cls)
+
+
+@pytest.mark.parametrize("truthy", ["true", "True", "TRUE"])
+def test_get_azure_key_vault_credential_uses_workload_identity_when_enabled(monkeypatch, truthy):
+    default_cls, workload_cls = _install_fake_azure_identity(monkeypatch)
+    monkeypatch.setenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY", truthy)
+
+    credential = ps._get_azure_key_vault_credential()
+
+    assert isinstance(credential, workload_cls)
+    assert not isinstance(credential, default_cls)
+
+
+@pytest.mark.parametrize("falsy", ["false", "False", "0", "unset-like"])
+def test_get_azure_key_vault_credential_uses_default_when_disabled(monkeypatch, falsy):
+    default_cls, workload_cls = _install_fake_azure_identity(monkeypatch)
+    monkeypatch.setenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY", falsy)
+
+    credential = ps._get_azure_key_vault_credential()
+
+    assert isinstance(credential, default_cls)
+
+
 # ---------------------------------------------------------------------------
 # cost_tracking
 # ---------------------------------------------------------------------------