From b40a7b5b536722f39734798b7983ab31872832d4 Mon Sep 17 00:00:00 2001 From: Stephen Sennett Date: Thu, 10 Sep 2026 01:13:55 +1000 Subject: [PATCH] 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 --- .../guardrail_hooks/azure/__init__.py | 11 +- .../guardrails/guardrail_hooks/azure/base.py | 53 ++++- .../guardrail_hooks/azure/prompt_shield.py | 9 +- .../guardrail_hooks/azure/text_moderation.py | 9 +- .../guardrails/guardrail_hooks/azure/base.py | 6 +- .../guardrail_hooks/azure/conftest.py | 44 ++++ .../guardrail_hooks/azure/test_azure_base.py | 188 ++++++++++++++++++ .../azure/test_azure_prompt_shield.py | 55 +++++ .../azure/test_azure_text_moderation.py | 23 +++ .../guardrails/test_guardrail_endpoints.py | 6 +- 10 files changed, 388 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py index de77cd55671..4f5b15df3d6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py +++ b/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, }, diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 42f0220cc4d..3989cc9934e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/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", } diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index 6e29d44662e..60247b9c4b1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 0dca8be3307..53bfac17910 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/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, ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py index c72888f33e1..861b1d2db02 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/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( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py new file mode 100644 index 00000000000..8e3605137ab --- /dev/null +++ b/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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py new file mode 100644 index 00000000000..3179be458e0 --- /dev/null +++ b/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" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py index 17e7222fa44..b421f52d3e0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py +++ b/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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index a43f95062f9..9df49db199d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/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" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 530f8ffd854..4d840394601 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/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