Merge 579e0d37f3 into 9071ca503e

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")
 
@@ -29,9 +32,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         ) = AzureContentSafetyPromptShieldGuardrail(
             guardrail_name=guardrail_name,
             **{
-                **litellm_params.model_dump(exclude_none=True),
+                **litellm_params.model_dump(exclude_unset=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,
             },
@@ -40,9 +44,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         azure_content_safety_guardrail = AzureContentSafetyTextModerationGuardrail(
             guardrail_name=guardrail_name,
             **{
-                **litellm_params.model_dump(exclude_none=True),
+                **litellm_params.model_dump(exclude_unset=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,5 +1,9 @@
+import asyncio
 import re
+from collections.abc import Callable
+from functools import cache
 from typing import TYPE_CHECKING, Any, Final
+from urllib.parse import urlparse
 
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -9,6 +13,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 +27,42 @@ 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"
+
+AZURE_CONTENT_SAFETY_ENTRA_HOST_SUFFIXES: Final = (
+    ".cognitiveservices.azure.com",
+    ".cognitiveservices.azure.us",
+    ".cognitiveservices.azure.cn",
+    ".services.ai.azure.com",
+)
+
+
+def _assert_entra_destination_is_azure(api_base: str) -> None:
+    """An Entra token is scoped to every Cognitive Services resource the identity can reach,
+    not to one resource, so it is only ever sent to an Azure endpoint over TLS. Entra also
+    requires the resource's custom subdomain, so any other host is not a valid destination."""
+    parsed: Final = urlparse(api_base)
+    host: Final = (parsed.hostname or "").lower()
+    if parsed.scheme == "https" and host.endswith(AZURE_CONTENT_SAFETY_ENTRA_HOST_SUFFIXES):
+        return
+    raise ValueError(
+        f"Azure Content Safety: refusing to send a Microsoft Entra token to api_base {api_base!r}. "
+        "Entra authentication needs the resource's HTTPS custom subdomain endpoint, for example "
+        "https://your-resource.cognitiveservices.azure.com. Set api_key to reach any other host"
+    )
+
+
+@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 +75,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 +89,33 @@ class AzureGuardrailBase:
         self.api_key = api_key
         self.api_base = api_base
         self.api_version: str = kwargs.get("api_version") or "2024-09-01"
+        if not api_key:
+            _assert_entra_destination_is_azure(api_base)
+        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)
+
+        _assert_entra_destination_is_azure(self.api_base)
+        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 +129,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,45 @@
+from collections.abc import Sequence
+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, Sequence[httpx.Request]]:
+    """An HTTP handler answering every Content Safety call, paired with the requests it saw."""
+    sent: Final[list[httpx.Request]] = []  # mutable-ok: callee-filled request log, handed back read-only
+
+    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,262 @@
+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]] = []  # mutable-ok: callee-filled thread log
+
+    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]] = []  # mutable-ok: callee-filled scope log
+
+    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"
+
+
+@pytest.mark.parametrize(
+    "bad_api_base",
+    [
+        "http://contoso.cognitiveservices.azure.com",
+        "https://contoso.cognitiveservices.azure.com.attacker.example",
+        "https://attacker.example",
+        "https://australiaeast.api.cognitive.microsoft.com",
+    ],
+)
+def test_keyless_guardrail_refuses_a_non_azure_destination(bad_api_base):
+    """The Entra token covers every Cognitive Services resource the identity can reach, so a
+    typo or a hijacked host would receive far more than one resource's key would give away."""
+    with pytest.raises(ValueError, match="refusing to send a Microsoft Entra token") as exc_info:
+        AzureContentSafetyPromptShieldGuardrail(guardrail_name="azure-guard", api_base=bad_api_base)
+
+    assert "api_key" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "good_api_base",
+    [
+        "https://contoso.cognitiveservices.azure.com",
+        "https://contoso.privatelink.cognitiveservices.azure.com",
+        "https://contoso.services.ai.azure.com",
+        "https://contoso.cognitiveservices.azure.us",
+    ],
+)
+def test_keyless_guardrail_accepts_azure_content_safety_endpoints(good_api_base):
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base=good_api_base,
+        entra_token_provider=lambda: "entra-token",
+    )
+
+    assert guardrail.api_base == good_api_base
+
+
+@pytest.mark.asyncio
+async def test_api_key_guardrail_may_use_any_destination(capturing_handler):
+    """A key is scoped to one resource, so gateways and test doubles stay reachable with one."""
+    handler, sent = capturing_handler
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base="https://gateway.internal.example",
+        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"
+
+
+@pytest.mark.asyncio
+async def test_clearing_api_key_cannot_send_a_token_to_a_non_azure_destination(capturing_handler):
+    """A guardrail admitted on its api_key must not start minting tokens for that same host."""
+    handler, sent = capturing_handler
+
+    guardrail: Final = AzureContentSafetyPromptShieldGuardrail(
+        guardrail_name="azure-guard",
+        api_base="https://gateway.internal.example",
+        api_key="secret-key",
+        entra_token_provider=lambda: "entra-token",
+    )
+    guardrail.async_handler = handler
+    guardrail.update_in_memory_litellm_params({"api_key": None})
+
+    with pytest.raises(ValueError, match="refusing to send a Microsoft Entra token"):
+        await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent == []

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,116 @@ 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
+
+
+@pytest.mark.asyncio
+async def test_config_without_api_version_uses_the_content_safety_default(api_base, capturing_handler):
+    """A config that omits api_version gets the Content Safety default, not another guardrail's."""
+    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",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent[0].url.params["api-version"] == "2024-09-01"
+
+
+@pytest.mark.asyncio
+async def test_config_api_version_is_honoured(api_base, capturing_handler):
+    """Pinning an older Content Safety version stays possible."""
+    handler, sent = capturing_handler
+
+    guardrail = initialize_guardrail(
+        LitellmParams(guardrail="azure/prompt_shield", mode="pre_call", api_base=api_base, api_version="2023-10-01"),
+        {"guardrail_name": "azure-prompt-shield"},
+        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].url.params["api-version"] == "2023-10-01"
+
+
+@pytest.mark.asyncio
+async def test_config_pricing_extras_survive_the_forwarded_params(api_base, capturing_handler):
+    """cost_tier and price_per_1000_text_records arrive as pydantic extras rather than declared
+    fields, so forwarding only the params the config set must still carry them through."""
+    handler, _ = capturing_handler
+
+    guardrail = initialize_guardrail(
+        LitellmParams(
+            guardrail="azure/prompt_shield",
+            mode="pre_call",
+            api_base=api_base,
+            cost_tier="paid",
+            price_per_1000_text_records=0.38,
+        ),
+        {"guardrail_name": "azure-prompt-shield"},
+        entra_token_provider=lambda: "entra-token",
+    )
+    guardrail.async_handler = handler
+    request_data = {"metadata": {}}
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data=request_data, input_type="request")
+
+    assert _recorded_guardrail_info(request_data)["guardrail_cost"] == pytest.approx(0.38 / 1000)

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

@@ -1,12 +1,15 @@
+import json
 from unittest.mock import Mock, patch
 
 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 +466,66 @@ 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"
+
+
+@pytest.mark.asyncio
+async def test_config_without_api_version_uses_the_content_safety_default(api_base, capturing_handler):
+    """A config that omits api_version gets the Content Safety default, not another guardrail's."""
+    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",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    assert sent[0].url.params["api-version"] == "2024-09-01"
+
+
+@pytest.mark.asyncio
+async def test_config_moderation_options_still_reach_the_request(api_base, capturing_handler):
+    """Forwarding only params the config set must not drop the guardrail's own options."""
+    handler, sent = capturing_handler
+
+    guardrail = initialize_guardrail(
+        LitellmParams(
+            guardrail="azure/text_moderations",
+            mode="pre_call",
+            api_base=api_base,
+            outputType="EightSeverityLevels",
+            blocklistNames=["my-blocklist"],
+        ),
+        {"guardrail_name": "azure-text-moderation"},
+        entra_token_provider=lambda: "entra-token",
+    )
+    guardrail.async_handler = handler
+
+    await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request")
+
+    body = json.loads(sent[0].content)
+    assert body["outputType"] == "EightSeverityLevels"
+    assert body["blocklistNames"] == ["my-blocklist"]

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