feat(guardrails): allow Entra ID auth for Azure Content Safety

Azure Content Safety accepts either an API key or a Microsoft Entra ID
bearer token, but the prompt shield and text moderation guardrails only
ever sent Ocp-Apim-Subscription-Key and refused to start without a key.
Organisations that disable local auth on the resource could not use them
at all

api_key is now optional. When it is omitted the guardrail authenticates
with the standard credential chain (env service principal, workload
identity, managed identity, az login) for the Cognitive Services scope
and sends Authorization: Bearer instead. api_base stays required, since
Entra auth needs the resource's custom subdomain endpoint

The credential is resolved once at startup, so a deployment missing
azure-identity fails at boot rather than on its first guarded request,
and a key-based deployment never imports azure-identity at all

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Final
 
 from litellm.types.guardrails import SupportedGuardrailIntegrations
@@ -9,11 +10,13 @@ if TYPE_CHECKING:
     from litellm.types.guardrails import Guardrail, LitellmParams
 
 
-def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
+def initialize_guardrail(
+    litellm_params: "LitellmParams",
+    guardrail: "Guardrail",
+    entra_token_provider: Callable[[], str] | None = None,
+) -> AzureContentSafetyPromptShieldGuardrail | AzureContentSafetyTextModerationGuardrail:
     import litellm
 
-    if not litellm_params.api_key:
-        raise ValueError("Azure Content Safety: api_key is required")
     if not litellm_params.api_base:
         raise ValueError("Azure Content Safety: api_base is required")
 
@@ -32,6 +35,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
                 **litellm_params.model_dump(exclude_none=True),
                 "api_key": litellm_params.api_key,
                 "api_base": litellm_params.api_base,
+                "entra_token_provider": entra_token_provider,
                 "default_on": litellm_params.default_on,
                 "event_hook": litellm_params.mode,
             },
@@ -43,6 +47,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
                 **litellm_params.model_dump(exclude_none=True),
                 "api_key": litellm_params.api_key,
                 "api_base": litellm_params.api_base,
+                "entra_token_provider": entra_token_provider,
                 "default_on": litellm_params.default_on,
                 "event_hook": litellm_params.mode,
             },

litellm/proxy/guardrails/guardrail_hooks/azure/base.py

@@ -1,4 +1,7 @@
+import asyncio
 import re
+from collections.abc import Callable
+from functools import cache
 from typing import TYPE_CHECKING, Any, Final
 
 from litellm._logging import verbose_proxy_logger
@@ -9,6 +12,9 @@ from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.secret_managers.get_azure_ad_token_provider import (
+    get_azure_ad_token_provider,
+)
 
 if TYPE_CHECKING:
     from litellm.types.llms.openai import AllMessageValues
@@ -20,6 +26,20 @@ AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000
 # chunk of N characters consumes ceil(N / 1000) text records.
 AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000
 
+AZURE_CONTENT_SAFETY_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default"
+
+
+@cache
+def _default_entra_token_provider() -> Callable[[], str]:
+    """Entra token provider for the Content Safety data plane, one credential per process."""
+    try:
+        return get_azure_ad_token_provider(azure_scope=AZURE_CONTENT_SAFETY_ENTRA_SCOPE)
+    except ImportError as e:
+        raise ValueError(
+            "Azure Content Safety: api_key is not set and azure-identity is not installed. "
+            "Set api_key, or install azure-identity to authenticate with Microsoft Entra ID"
+        ) from e
+
 
 class AzureGuardrailBase:
     """
@@ -32,8 +52,10 @@ class AzureGuardrailBase:
 
     def __init__(
         self,
-        api_key: str,
+        *,
         api_base: str,
+        api_key: str | None = None,
+        entra_token_provider: Callable[[], str] | None = None,
         **kwargs: Any,
     ):
         # Forward remaining kwargs to the next class in the MRO
@@ -44,6 +66,30 @@ class AzureGuardrailBase:
         self.api_key = api_key
         self.api_base = api_base
         self.api_version: str = kwargs.get("api_version") or "2024-09-01"
+        self._entra_token_provider: Final = entra_token_provider or (
+            None if api_key else _default_entra_token_provider()
+        )
+
+    async def _auth_header(self) -> tuple[str, str]:
+        """Credential header name and value for a single request.
+
+        Azure Content Safety accepts an API key or a Microsoft Entra token and rejects each
+        on the other's header, so exactly one is sent.
+        """
+        if self.api_key:
+            return ("Ocp-Apim-Subscription-Key", self.api_key)
+
+        minter: Final = self._entra_token_provider or _default_entra_token_provider()
+        try:
+            token: Final = await asyncio.to_thread(minter)
+        except Exception as e:
+            verbose_proxy_logger.exception("Azure Content Safety: Entra token request failed")
+            raise ValueError(
+                "Azure Content Safety: no credential available. Set api_key, or configure an Entra "
+                "identity (AZURE_CLIENT_ID / AZURE_CLIENT_SECRET / AZURE_TENANT_ID, workload identity, "
+                "managed identity, or az login) holding the Cognitive Services User role on the resource"
+            ) from e
+        return ("Authorization", f"Bearer {token}")
 
     async def _post_to_content_safety(self, endpoint_path: str, request_body: dict[str, object]) -> dict[str, Any]:
         """POST to an Azure Content Safety endpoint with standard auth headers.
@@ -57,8 +103,9 @@ class AzureGuardrailBase:
             Parsed JSON response dict.
         """
         url: Final = f"{self.api_base}/contentsafety/{endpoint_path}?api-version={self.api_version}"
-        headers: Final = {
-            "Ocp-Apim-Subscription-Key": self.api_key,
+        auth_name, auth_value = await self._auth_header()
+        headers: Final = {  # mutable-ok: AsyncHTTPHandler.post types headers as `dict | None`
+            auth_name: auth_value,
             "Content-Type": "application/json",
         }
 

litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py

@@ -120,8 +120,8 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
 
     Configuration:
         guardrail_name: Name of the guardrail instance
-        api_key: Azure Prompt Shield API key
-        api_base: Azure Prompt Shield API endpoint
+        api_base: Azure Content Safety endpoint
+        api_key: Azure Content Safety API key. Omit it to authenticate with Microsoft Entra ID
         default_on: Whether to enable by default
     """
 
@@ -129,17 +129,18 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
 
     def __init__(
         self,
+        *,
         guardrail_name: str,
-        api_key: str,
         api_base: str,
+        api_key: str | None = None,
         **kwargs,
     ):
         """Initialize Azure Prompt Shield guardrail handler."""
         # AzureGuardrailBase.__init__ stores api_key, api_base, api_version,
         # async_handler and forwards the rest to CustomGuardrail.
         super().__init__(
-            api_key=api_key,
             api_base=api_base,
+            api_key=api_key,
             guardrail_name=guardrail_name,
             supported_event_hooks=list(self.get_supported_event_hooks()),
             **kwargs,

litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py

@@ -37,8 +37,8 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
 
     Configuration:
         guardrail_name: Name of the guardrail instance
-        api_key: Azure Text Moderation API key
-        api_base: Azure Text Moderation API endpoint
+        api_base: Azure Content Safety endpoint
+        api_key: Azure Content Safety API key. Omit it to authenticate with Microsoft Entra ID
         default_on: Whether to enable by default
     """
 
@@ -55,9 +55,10 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
 
     def __init__(
         self,
+        *,
         guardrail_name: str,
-        api_key: str,
         api_base: str,
+        api_key: str | None = None,
         severity_threshold: int | None = None,
         severity_threshold_by_category: dict[str, int] | None = None,
         **kwargs,
@@ -71,8 +72,8 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
         # AzureGuardrailBase.__init__ stores api_key, api_base, api_version,
         # async_handler and forwards the rest to CustomGuardrail.
         super().__init__(
-            api_key=api_key,
             api_base=api_base,
+            api_key=api_key,
             guardrail_name=guardrail_name,
             **kwargs,
         )

litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py

@@ -6,7 +6,11 @@ class AzureContentSafetyConfigModel(BaseModel):
 
     api_key: str | None = Field(
         default=None,
-        description="API key for the Azure Content Safety Prompt Shield guardrail",
+        description=(
+            "API key for the Azure Content Safety resource. Optional: omit it to authenticate with "
+            "Microsoft Entra ID, which needs the Cognitive Services User role on the resource and a "
+            "custom subdomain api_base"
+        ),
     )
 
     api_base: str | None = Field(

tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py

@@ -0,0 +1,44 @@
+from typing import Final
+
+import httpx
+import pytest
+
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.proxy.guardrails.guardrail_hooks.azure.base import (
+    _default_entra_token_provider,
+)
+
+CLEAN_RESPONSE_FOR_BOTH_GUARDRAILS: Final = {
+    "userPromptAnalysis": {"attackDetected": False},
+    "documentsAnalysis": [],
+    "categoriesAnalysis": [],
+    "blocklistsMatch": [],
+}
+
+
+@pytest.fixture(autouse=True)
+def clear_default_entra_provider_cache():
+    """A credential resolved elsewhere in the session would otherwise be reused here, masking
+    the dependency and failure paths these tests assert."""
+    _default_entra_token_provider.cache_clear()
+    yield
+    _default_entra_token_provider.cache_clear()
+
+
+@pytest.fixture
+def api_base() -> str:
+    return "https://contoso.cognitiveservices.azure.com"
+
+
+@pytest.fixture
+def capturing_handler() -> tuple[AsyncHTTPHandler, list[httpx.Request]]:
+    """An HTTP handler answering every Content Safety call, paired with the requests it saw."""
+    sent: Final[list[httpx.Request]] = []
+
+    def _record(request: httpx.Request) -> httpx.Response:
+        sent.append(request)
+        return httpx.Response(200, json=CLEAN_RESPONSE_FOR_BOTH_GUARDRAILS)
+
+    handler: Final = AsyncHTTPHandler()
+    handler.client = httpx.AsyncClient(transport=httpx.MockTransport(_record))
+    return handler, sent

tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py

@@ -0,0 +1,188 @@
+import sys
+import threading
+from typing import Final
+
+import pytest
+
+from litellm.proxy.guardrails.guardrail_hooks.azure import base
+from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import (
+    AzureContentSafetyPromptShieldGuardrail,
+)
+from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import (
+    AzureContentSafetyTextModerationGuardrail,
+)
+
+GUARDRAIL_CLASSES: Final = (
+    AzureContentSafetyPromptShieldGuardrail,
+    AzureContentSafetyTextModerationGuardrail,
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("guardrail_class", GUARDRAIL_CLASSES)
+async def test_api_key_rides_the_subscription_key_header_and_mints_no_token(
+    guardrail_class, api_base, capturing_handler
+):
+    handler, sent = capturing_handler
+
+    def _must_not_mint() -> str:
+        raise AssertionError("Entra token minted despite api_key being set")
+
+    guardrail: Final = guardrail_class(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        api_key="secret-key",
+        entra_token_provider=_must_not_mint,
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert len(sent) == 1
+    assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key"
+    assert "authorization" not in sent[0].headers
+    assert sent[0].headers["Content-Type"] == "application/json"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("guardrail_class", GUARDRAIL_CLASSES)
+async def test_omitted_api_key_rides_an_entra_bearer_token(guardrail_class, api_base, capturing_handler):
+    handler, sent = capturing_handler
+
+    guardrail: Final = guardrail_class(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        entra_token_provider=lambda: "entra-token",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert len(sent) == 1
+    assert sent[0].headers["Authorization"] == "Bearer entra-token"
+    assert "ocp-apim-subscription-key" not in sent[0].headers
+    assert sent[0].headers["Content-Type"] == "application/json"
+
+
+@pytest.mark.asyncio
+async def test_blank_api_key_is_treated_as_absent(api_base, capturing_handler):
+    """An `os.environ/` api_key resolving to an empty string must reach Entra, not send a blank key."""
+    handler, sent = capturing_handler
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        api_key="",
+        entra_token_provider=lambda: "entra-token",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent[0].headers["Authorization"] == "Bearer entra-token"
+    assert "ocp-apim-subscription-key" not in sent[0].headers
+
+
+@pytest.mark.asyncio
+async def test_token_minting_failure_reports_the_options_and_sends_nothing(api_base, capturing_handler):
+    handler, sent = capturing_handler
+
+    def _fails() -> str:
+        raise RuntimeError("no managed identity endpoint found")
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        entra_token_provider=_fails,
+    )
+    guardrail.async_handler = handler
+
+    with pytest.raises(ValueError, match="no credential available") as exc_info:
+        await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent == []
+    assert isinstance(exc_info.value.__cause__, RuntimeError)
+
+
+@pytest.mark.asyncio
+async def test_token_minting_failure_keeps_credential_detail_out_of_the_error(api_base, capturing_handler):
+    """azure-identity errors carry tenant and client ids, and this message reaches the API caller."""
+    handler, _ = capturing_handler
+
+    def _fails() -> str:
+        raise RuntimeError("tenant 11111111-2222-3333-4444-555555555555 rejected the request")
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        entra_token_provider=_fails,
+    )
+    guardrail.async_handler = handler
+
+    with pytest.raises(ValueError, match="no credential available") as exc_info:
+        await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert "11111111-2222-3333-4444-555555555555" not in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_token_is_minted_off_the_event_loop_thread(api_base, capturing_handler):
+    """Credential sources block: IMDS probes time out and the az CLI credential spawns a subprocess."""
+    handler, _ = capturing_handler
+    minting_threads: Final[list[int]] = []
+
+    def _record_thread() -> str:
+        minting_threads.append(threading.get_ident())
+        return "entra-token"
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        entra_token_provider=_record_thread,
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert len(minting_threads) == 1
+    assert minting_threads[0] != threading.get_ident()
+
+
+def test_default_credential_is_built_once_per_process(monkeypatch):
+    """Rebuilding it per request costs a fresh credential and token round trip on every scan."""
+    builds: Final[list[str]] = []
+
+    def _build(azure_scope: str):
+        builds.append(azure_scope)
+        return lambda: "entra-token"
+
+    monkeypatch.setattr(base, "get_azure_ad_token_provider", _build)
+
+    assert base._default_entra_token_provider() is base._default_entra_token_provider()
+    assert builds == [base.AZURE_CONTENT_SAFETY_ENTRA_SCOPE]
+
+
+@pytest.mark.parametrize("guardrail_class", GUARDRAIL_CLASSES)
+def test_keyless_guardrail_names_the_missing_dependency_at_startup(guardrail_class, api_base, monkeypatch):
+    monkeypatch.setitem(sys.modules, "azure", None)
+
+    with pytest.raises(ValueError, match="azure-identity"):
+        guardrail_class(guardrail_name="azure-guard", api_base=api_base)
+
+
+@pytest.mark.asyncio
+async def test_api_key_guardrail_never_reaches_azure_identity(api_base, capturing_handler, monkeypatch):
+    """azure-identity ships in the proxy extra, so a key-based deployment must not depend on it."""
+    handler, sent = capturing_handler
+    monkeypatch.setitem(sys.modules, "azure", None)
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base=api_base,
+        api_key="secret-key",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key"

tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py

@@ -4,6 +4,7 @@ import pytest
 from fastapi import HTTPException
 
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.azure import initialize_guardrail
 from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import (
     AzureContentSafetyPromptShieldGuardrail,
 )
@@ -635,3 +636,57 @@ def test_update_in_memory_litellm_params_dead_env_credential_rejected_untouched(
 
     assert guardrail.api_key == "azure_prompt_shield_api_key"
     assert guardrail.price_per_1000_text_records == 0.38
+
+
+@pytest.mark.asyncio
+async def test_initialize_guardrail_without_api_key_authenticates_with_entra(api_base, capturing_handler):
+    """A keyless config entry yields a guardrail that authenticates with Entra."""
+    handler, sent = capturing_handler
+
+    guardrail = initialize_guardrail(
+        LitellmParams(guardrail="azure/prompt_shield", mode="pre_call", api_base=api_base),
+        {"guardrail_name": "azure-prompt-shield"},
+        entra_token_provider=lambda: "entra-token",
+    )
+
+    assert isinstance(guardrail, AzureContentSafetyPromptShieldGuardrail)
+    assert guardrail.api_key is None
+    assert guardrail.api_base == api_base
+
+    guardrail.async_handler = handler
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent[0].headers["Authorization"] == "Bearer entra-token"
+
+
+def test_initialize_guardrail_without_api_base_still_raises():
+    """api_base carries the resource's custom subdomain, which Entra auth cannot work without."""
+    with pytest.raises(ValueError, match="api_base is required"):
+        initialize_guardrail(
+            LitellmParams(guardrail="azure/prompt_shield", mode="pre_call"),
+            {"guardrail_name": "azure-prompt-shield"},
+            entra_token_provider=lambda: "entra-token",
+        )
+
+
+@pytest.mark.asyncio
+async def test_clearing_api_key_at_runtime_switches_to_entra(api_base, capturing_handler):
+    """A dashboard edit that removes the key must re-authenticate, not keep sending a stale header."""
+    handler, sent = capturing_handler
+
+    guardrail = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure_prompt_shield",
+        api_base=api_base,
+        api_key="secret-key",
+        entra_token_provider=lambda: "entra-token",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+    guardrail.update_in_memory_litellm_params({"api_key": None})
+    await guardrail.apply_guardrail(inputs={"texts": ["hello again"]}, request_data={}, input_type="request")
+
+    assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key"
+    assert "authorization" not in sent[0].headers
+    assert sent[1].headers["Authorization"] == "Bearer entra-token"
+    assert "ocp-apim-subscription-key" not in sent[1].headers

tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py

@@ -4,9 +4,11 @@ import pytest
 from fastapi import HTTPException
 
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.azure import initialize_guardrail
 from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import (
     AzureContentSafetyTextModerationGuardrail,
 )
+from litellm.types.guardrails import LitellmParams
 from litellm.types.utils import Choices, Message, ModelResponse
 
 
@@ -463,3 +465,24 @@ async def test_apply_guardrail_handles_missing_texts_key():
 
     mock_post.assert_not_called()
     assert result == {"images": ["x"]}
+
+
+@pytest.mark.asyncio
+async def test_initialize_guardrail_without_api_key_authenticates_with_entra(api_base, capturing_handler):
+    """A keyless config entry yields a guardrail that authenticates with Entra."""
+    handler, sent = capturing_handler
+
+    guardrail = initialize_guardrail(
+        LitellmParams(guardrail="azure/text_moderations", mode="pre_call", api_base=api_base),
+        {"guardrail_name": "azure-text-moderation"},
+        entra_token_provider=lambda: "entra-token",
+    )
+
+    assert isinstance(guardrail, AzureContentSafetyTextModerationGuardrail)
+    assert guardrail.api_key is None
+    assert guardrail.api_base == api_base
+
+    guardrail.async_handler = handler
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent[0].headers["Authorization"] == "Bearer entra-token"

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -632,7 +632,11 @@ def test_get_provider_specific_params():
     # Check the structure of a simple field
     assert (
         fields["api_key"]["description"]
-        == "API key for the Azure Content Safety Prompt Shield guardrail"
+        == (
+            "API key for the Azure Content Safety resource. Optional: omit it to authenticate with "
+            "Microsoft Entra ID, which needs the Cognitive Services User role on the resource and a "
+            "custom subdomain api_base"
+        )
     )
     assert fields["api_key"]["required"] == False
     assert fields["api_key"]["type"] == "string"  # Should be string, not None