diff --git a/.circleci/config.yml b/.circleci/config.yml index 38fdaf3609d..cf69ff68da6 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2541,7 +2541,6 @@ jobs: paths: - litellm-docker-database.tar.zst - test_bad_database_url: machine: image: ubuntu-2204:2024.04.1 diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 3ee56dfc5ca..32ae61d7f58 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1204,12 +1204,8 @@ def get_last_user_message(messages: List[AllMessageValues]) -> Optional[str]: {"role": "assistant", "content": "I'm good, thank you!"}, {"role": "user", "content": "What is the weather in Tokyo?"}, ] - get_user_prompt(messages) -> "What is the weather in Tokyo?" + get_last_user_message(messages) -> "What is the weather in Tokyo?" """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - convert_content_list_to_str, - ) - if not messages: return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index fab65884a9f..7d2dfce0711 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -2,6 +2,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_last_user_message, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -134,32 +137,4 @@ def get_user_prompt(self, messages: List["AllMessageValues"]) -> Optional[str]: ] get_user_prompt(messages) -> "What is the weather in Tokyo?" """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - convert_content_list_to_str, - ) - - if not messages: - return None - - # Iterate from the end to find the last consecutive block of user messages - user_messages = [] - for message in reversed(messages): - if message.get("role") == "user": - user_messages.append(message) - else: - # Stop when we hit a non-user message - break - - if not user_messages: - return None - - # Reverse to get the messages in chronological order - user_messages.reverse() - - user_prompt = "" - for message in user_messages: - text_content = convert_content_list_to_str(message) - user_prompt += text_content + "\n" - - result = user_prompt.strip() - return result if result else None + return get_last_user_message(messages) diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py new file mode 100644 index 00000000000..9ead2a63b60 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py @@ -0,0 +1,57 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .purview_dlp import MicrosoftPurviewDLPGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + tenant_id = getattr(litellm_params, "tenant_id", None) + client_id = getattr(litellm_params, "client_id", None) + + # client_secret can be passed via the standard api_key field or as + # a dedicated client_secret parameter. + client_secret = litellm_params.api_key or getattr( + litellm_params, "client_secret", None + ) + + if not tenant_id: + raise ValueError("Microsoft Purview: tenant_id is required") + if not client_id: + raise ValueError("Microsoft Purview: client_id is required") + if not client_secret: + raise ValueError("Microsoft Purview: client_secret (or api_key) is required") + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Microsoft Purview: guardrail_name is required") + + purview_guardrail = MicrosoftPurviewDLPGuardrail( + guardrail_name=guardrail_name, + tenant_id=str(tenant_id), + client_id=str(client_id), + client_secret=str(client_secret), + purview_app_name=str( + getattr(litellm_params, "purview_app_name", None) or "LiteLLM" + ), + user_id_field=str(getattr(litellm_params, "user_id_field", None) or "user_id"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(purview_guardrail) + return purview_guardrail + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.MICROSOFT_PURVIEW.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.MICROSOFT_PURVIEW.value: MicrosoftPurviewDLPGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py new file mode 100644 index 00000000000..a7ed1d40913 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -0,0 +1,515 @@ +import threading +import time +import uuid +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) + +if TYPE_CHECKING: + from litellm.types.llms.openai import AllMessageValues + +GRAPH_API_BASE = "https://graph.microsoft.com/v1.0" +TOKEN_ENDPOINT_TEMPLATE = ( + "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token" +) +GRAPH_SCOPE = "https://graph.microsoft.com/.default" + +# Protection scope cache TTL in seconds (1 hour, per Microsoft recommendation). +SCOPE_CACHE_TTL_SECONDS = 3600.0 + + +class PurviewGuardrailBase: + """ + Base class for Microsoft Purview guardrails. + + Manages OAuth2 client-credentials token acquisition, protection scope + computation with ETag caching, and authenticated POST calls to the + Microsoft Graph API. + """ + + def __init__( + self, + tenant_id: str, + client_id: str, + client_secret: str, + purview_app_name: str = "LiteLLM", + user_id_field: str = "user_id", + **kwargs: Any, + ): + # Forward remaining kwargs to the next class in the MRO + # (typically CustomGuardrail). + super().__init__(**kwargs) + + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + self.tenant_id = tenant_id + self.client_id = client_id + self.client_secret = client_secret + self.purview_app_name = purview_app_name + self.user_id_field = user_id_field + + # Token cache: (access_token, expires_at_epoch) + self._token_cache: Optional[Tuple[str, float]] = None + + # Protection scope cache: user_id -> (etag, scope_response, fetched_at) + # Capped at 1000 entries (LRU eviction) to avoid unbounded growth. + self._scope_cache: OrderedDict[str, Tuple[str, Dict[str, Any], float]] = ( + OrderedDict() + ) + self._scope_cache_maxsize = 1000 + # Use a threading.Lock (not asyncio.Lock) because this lock is acquired + # from both the proxy's main asyncio event loop and from short-lived + # event loops created by the logging_hook thread fallback. In Python + # 3.10+ an asyncio.Lock is bound to the first event loop that acquires + # it and raises RuntimeError from any other loop, which would silently + # break audit logging via the thread fallback. All critical sections + # below are pure in-memory dict ops with no awaits, so a synchronous + # lock is both correct and sufficient. + self._cache_lock = threading.Lock() + + @staticmethod + def _encode_graph_user_id(user_id: str) -> str: + """Percent-encode Entra user id for Graph ``/users/{id}/...`` path segments.""" + return encode_url_path_segment(user_id, field_name="user_id") + + # ------------------------------------------------------------------ + # OAuth2 token management + # ------------------------------------------------------------------ + + async def _get_access_token(self) -> str: + """Acquire or return cached OAuth2 token via client_credentials grant.""" + now = time.time() + with self._cache_lock: + if self._token_cache and self._token_cache[1] > now + 60: + return self._token_cache[0] + + url = TOKEN_ENDPOINT_TEMPLATE.format(tenant_id=self.tenant_id) + data = { + "grant_type": "client_credentials", + "client_id": self.client_id, + "client_secret": self.client_secret, + "scope": GRAPH_SCOPE, + } + response = await self.async_handler.post( + url=url, + data=data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + response.raise_for_status() + token_data = response.json() + access_token = token_data["access_token"] + expires_in = int(token_data.get("expires_in", 3599)) + # Recompute ``now`` after the await so the expiry reflects when the + # token was actually received, not when the request started. + with self._cache_lock: + self._token_cache = (access_token, time.time() + expires_in) + verbose_proxy_logger.debug( + "Purview: acquired new OAuth2 token (expires_in=%ds)", expires_in + ) + return access_token + + # ------------------------------------------------------------------ + # Graph API helpers + # ------------------------------------------------------------------ + + async def _graph_post( + self, + url: str, + json_body: Dict[str, Any], + extra_headers: Optional[Dict[str, str]] = None, + ) -> Tuple[Dict[str, Any], Dict[str, str]]: + """POST to Graph API with bearer auth. + + Returns: + Tuple of (response_json, response_headers). + """ + token = await self._get_access_token() + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + if extra_headers: + headers.update(extra_headers) + + verbose_proxy_logger.debug("Purview Graph POST %s", url) + response = await self.async_handler.post( + url=url, headers=headers, json=json_body + ) + response.raise_for_status() + response_json: Dict[str, Any] = response.json() + response_headers = dict(response.headers) + verbose_proxy_logger.debug("Purview Graph response: %s", response_json) + return response_json, response_headers + + # ------------------------------------------------------------------ + # Protection scopes + # ------------------------------------------------------------------ + + async def _compute_protection_scopes( + self, user_id: str + ) -> Tuple[str, Dict[str, Any]]: + """Call protectionScopes/compute and cache with ETag. + + Returns: + Tuple of (etag, scope_response). + """ + encoded_user_id = self._encode_graph_user_id(user_id) + now = time.time() + + with self._cache_lock: + cached = self._scope_cache.get(user_id) + if cached and (now - cached[2]) < SCOPE_CACHE_TTL_SECONDS: + self._scope_cache.move_to_end(user_id) + return cached[0], cached[1] + + url = ( + f"{GRAPH_API_BASE}/users/{encoded_user_id}" + "/dataSecurityAndGovernance/protectionScopes/compute" + ) + body: Dict[str, Any] = { + "activities": "uploadText,downloadText", + "locations": [ + { + "@odata.type": "microsoft.graph.policyLocationApplication", + "value": self.client_id, + } + ], + } + + response_json, response_headers = await self._graph_post(url, body) + etag = response_headers.get("etag", response_headers.get("ETag", "")) + + # Recompute ``now`` after the await so the TTL reflects when the + # scope response was actually received, not when the request started. + fetched_at = time.time() + with self._cache_lock: + self._scope_cache[user_id] = (etag, response_json, fetched_at) + # Move refreshed entry to the end so it is treated as most-recently-used. + # OrderedDict.__setitem__ preserves existing insertion order for known + # keys, so an explicit move_to_end() call is required. + self._scope_cache.move_to_end(user_id) + # Evict least-recently-used entry when cache exceeds max size. + while len(self._scope_cache) > self._scope_cache_maxsize: + self._scope_cache.popitem(last=False) + return etag, response_json + + # ------------------------------------------------------------------ + # Process content + # ------------------------------------------------------------------ + + async def _process_content( + self, + user_id: str, + text: str, + activity: str, + etag: str, + correlation_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Call processContent for DLP policy evaluation. + + Args: + user_id: Entra object ID of the user. + text: The content to evaluate. + activity: ``"uploadText"`` for prompts, ``"downloadText"`` for responses. + etag: Cached ETag from protectionScopes/compute. + correlation_id: Optional conversation/thread ID. + """ + encoded_user_id = self._encode_graph_user_id(user_id) + url = ( + f"{GRAPH_API_BASE}/users/{encoded_user_id}" + "/dataSecurityAndGovernance/processContent" + ) + body: Dict[str, Any] = { + "contentToProcess": { + "contentEntries": [ + { + "@odata.type": "microsoft.graph.processConversationMetadata", + "identifier": str(uuid.uuid4()), + "content": { + "@odata.type": "microsoft.graph.textContent", + "data": text, + }, + "name": f"{self.purview_app_name} message", + "correlationId": correlation_id or str(uuid.uuid4()), + "sequenceNumber": 0, + "isTruncated": False, + } + ], + "activityMetadata": {"activity": activity}, + "deviceMetadata": {}, + "protectedAppMetadata": { + "name": self.purview_app_name, + "version": "1.0", + "applicationLocation": { + "@odata.type": "microsoft.graph.policyLocationApplication", + "value": self.client_id, + }, + }, + "integratedAppMetadata": { + "name": self.purview_app_name, + "version": "1.0", + }, + } + } + + extra_headers: Dict[str, str] = {} + if etag: + extra_headers["If-None-Match"] = etag + + response_json, _ = await self._graph_post(url, body, extra_headers) + + # If policies changed, invalidate scope cache so next call re-fetches. + if response_json.get("protectionScopeState") == "modified": + with self._cache_lock: + self._scope_cache.pop(user_id, None) + + return response_json + + # ------------------------------------------------------------------ + # User ID resolution + # ------------------------------------------------------------------ + + def _resolve_user_id( + self, data: Dict[str, Any], user_api_key_dict: Any + ) -> Optional[str]: + """Resolve the Entra user object ID from request data or auth context. + + Returns the strongest available identity walking down four sources, in + decreasing trust order: + + 1. ``user_api_key_dict.user_id`` — LiteLLM key / JWT-bound user + 2. ``user_api_key_dict.end_user_id`` — request-derived + 3. ``metadata["user_api_key_user_id"]`` — proxy-injected from the key + 4. ``metadata[user_id_field]`` — caller-supplied + + Used only by blocking-mode resolution to disambiguate "no identity at + all" from "caller supplied an untrusted identity" for the error + message. Neither blocking nor audit DLP feeds the untrusted + fallbacks (2, 4) into Purview itself. + """ + trusted = self._resolve_trusted_user_id(data, user_api_key_dict) + if trusted: + return trusted + + if hasattr(user_api_key_dict, "end_user_id") and user_api_key_dict.end_user_id: + return str(user_api_key_dict.end_user_id) + + metadata = data.get("metadata") or data.get("litellm_metadata") or {} + uid = metadata.get("user_api_key_user_id") + if uid: + return str(uid) + + uid = metadata.get(self.user_id_field) + if uid: + return str(uid) + + return None + + @staticmethod + def _logging_kwargs_metadata(kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Metadata dict from ``model_call_details`` / logging kwargs.""" + litellm_params = kwargs.get("litellm_params") or {} + if not isinstance(litellm_params, dict): + return {} + md = litellm_params.get("metadata") + return md if isinstance(md, dict) else {} + + def _resolve_trusted_user_id( + self, data: Dict[str, Any], user_api_key_dict: Any + ) -> Optional[str]: + """Resolve user ID from API-key/JWT-bound identity for blocking DLP. + + Uses only ``UserAPIKeyAuth.user_id`` (bound on the LiteLLM key or JWT). + Intentionally omits ``UserAPIKeyAuth.end_user_id`` because the proxy sets + it from caller-controlled request fields (``user``, ``metadata.user_id``, + ``safety_identifier``, custom headers, etc.) via + ``get_end_user_id_from_request_body``. + + Also omits ``metadata[user_id_field]`` and + ``metadata["user_api_key_user_id"]`` for the same impersonation risk when + the key has no bound user. + + Returns ``None`` when no authenticated identity is available. Blocking + hooks must fail closed rather than skip the DLP check. + """ + if hasattr(user_api_key_dict, "user_id") and user_api_key_dict.user_id: + return str(user_api_key_dict.user_id) + + return None + + def _resolve_user_id_from_logging_kwargs( + self, kwargs: Dict[str, Any] + ) -> Optional[str]: + """Trusted-identity-only resolver for logging-only hooks. + + Uses only the proxy-injected ``user_api_key_user_id`` (populated from + the API-key/JWT-bound ``UserAPIKeyAuth.user_id`` after the proxy + strips every caller-supplied ``user_api_key_*`` key from the request + metadata). Caller-influenceable sources (``user_api_key_end_user_id``, + ``metadata[user_id_field]``) are not used here so a caller cannot + cause Purview audit records to be written under a victim's identity. + Returns ``None`` when no trusted identity is available so the audit + is skipped rather than misattributed. + """ + md = self._logging_kwargs_metadata(kwargs) + uid = md.get("user_api_key_user_id") or kwargs.get("user_api_key_user_id") + if uid: + return str(uid) + return None + + # ------------------------------------------------------------------ + # Policy action evaluation + # ------------------------------------------------------------------ + + @staticmethod + def _should_block(response: Dict[str, Any]) -> bool: + """Return True if any policyAction requires blocking.""" + for action in response.get("policyActions", []): + odata_type = action.get("@odata.type", "") + action_field = action.get("action", "") + + if "restrictAccessAction" in odata_type or action_field == "restrictAccess": + restriction = action.get("restrictionAction", "") + if restriction == "block": + return True + return False + + # ------------------------------------------------------------------ + # Prompt text for DLP + # ------------------------------------------------------------------ + + @staticmethod + def is_token_id_prompt(prompt: Any) -> bool: + """Return True if ``prompt`` carries OpenAI completions token ids. + + Covers every list shape that ``completion_prompt_to_str`` cannot decode + for Purview, including flat ``list[int]`` (single token-id prompt), + ``list[list[int]]`` (multi-prompt token-id batches), and mixed lists + that include any token-id sub-array. + """ + if not isinstance(prompt, list) or not prompt: + return False + for x in prompt: + if isinstance(x, int): + return True + if isinstance(x, list) and x and any(isinstance(y, int) for y in x): + return True + return False + + @staticmethod + def completion_prompt_to_str(prompt: Any) -> Optional[str]: + """Normalize OpenAI ``/v1/completions`` ``prompt`` for text DLP. + + Supports string prompts and list-of-string prompts. List-of-token-id prompts + are skipped (no plaintext for Purview to evaluate). + """ + if prompt is None: + return None + if isinstance(prompt, str): + stripped = prompt.strip() + return stripped or None + if isinstance(prompt, list) and prompt: + if all(isinstance(x, str) for x in prompt): + joined = "\n".join(s.strip() for s in prompt if isinstance(s, str)) + return joined.strip() or None + if all(isinstance(x, int) for x in prompt): + verbose_proxy_logger.debug( + "Purview DLP: completions prompt is token ids only; skipping text scan" + ) + return None + str_parts = [x for x in prompt if isinstance(x, str)] + if str_parts: + joined = "\n".join(s.strip() for s in str_parts) + return joined.strip() or None + return None + + @staticmethod + def _extract_tool_call_args_from_message(message: Any) -> List[str]: + """Return plaintext arguments strings from tool_calls and function_call fields. + + Covers both the request path (assistant messages in chat histories that + carry tool_calls / function_call) and the response path (model-generated + tool calls returned in a ModelResponse). Both dict-style and object-style + representations are handled. + """ + args: List[str] = [] + + # tool_calls: [{"function": {"arguments": "..."}}] + tool_calls = ( + message.get("tool_calls") + if isinstance(message, dict) + else getattr(message, "tool_calls", None) + ) + if tool_calls: + for tc in tool_calls: + fn = ( + tc.get("function") + if isinstance(tc, dict) + else getattr(tc, "function", None) + ) + if fn is None: + continue + arguments = ( + fn.get("arguments") + if isinstance(fn, dict) + else getattr(fn, "arguments", None) + ) + if isinstance(arguments, str) and arguments.strip(): + args.append(arguments) + + # Legacy function_call: {"arguments": "..."} + function_call = ( + message.get("function_call") + if isinstance(message, dict) + else getattr(message, "function_call", None) + ) + if function_call is not None: + arguments = ( + function_call.get("arguments") + if isinstance(function_call, dict) + else getattr(function_call, "arguments", None) + ) + if isinstance(arguments, str) and arguments.strip(): + args.append(arguments) + + return args + + def get_prompt_text_for_dlp( + self, messages: List["AllMessageValues"] + ) -> Optional[str]: + """Concatenate text from every chat message (all roles) for pre-call DLP. + + Evaluates the same payload the model receives, not only the trailing user + turn. Each message is separated by ``\\n\\n`` so that tokens at message + boundaries are not merged (e.g., ``"end of msg1\\n\\nstart of msg2"`` + rather than ``"end of msg1start of msg2"``), which preserves DLP pattern + detection accuracy across message boundaries. + + Tool-call arguments (``tool_calls[].function.arguments`` and + ``function_call.arguments``) are included alongside message content so + that sensitive data hidden in function arguments is not bypassed. + """ + if not messages: + return None + parts: List[str] = [] + for msg in messages: + segments: List[str] = [] + content = convert_content_list_to_str(message=msg).strip() + if content: + segments.append(content) + segments.extend(self._extract_tool_call_args_from_message(msg)) + combined = "\n".join(segments) + if combined.strip(): + parts.append(combined.strip()) + text = "\n\n".join(parts) + return text or None diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py new file mode 100644 index 00000000000..ee0bac64d4f --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py @@ -0,0 +1,734 @@ +""" +Microsoft Purview DLP Guardrail for LiteLLM. + +Supports three modes: +- pre_call: Block sensitive data in prompts before they reach the LLM. +- post_call: Block sensitive data in LLM responses. +- logging_only: Log interactions to Purview for audit/compliance without blocking. +""" + +import asyncio +import threading +import uuid +from datetime import datetime +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Dict, + List, + Optional, + Tuple, + Type, + Union, + cast, +) + +import httpx +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + Choices, + GuardrailStatus, + ModelResponse, + ModelResponseStream, + ResponsesAPIResponse, + TextChoices, + TextCompletionResponse, +) + +from .base import PurviewGuardrailBase + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.proxy.guardrails.guardrail_hooks.base import ( + GuardrailConfigModel, + ) + from litellm.types.utils import ( + CallTypesLiteral, + EmbeddingResponse, + ImageResponse, + ) + + +class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): + """ + Microsoft Purview DLP guardrail. + + Evaluates prompts and responses against Microsoft Purview DLP policies + via the Microsoft Graph ``processContent`` API. + """ + + def __init__( + self, + guardrail_name: str, + tenant_id: str, + client_id: str, + client_secret: str, + purview_app_name: str = "LiteLLM", + user_id_field: str = "user_id", + **kwargs: Any, + ): + supported_event_hooks = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.logging_only, + ] + + super().__init__( + tenant_id=tenant_id, + client_id=client_id, + client_secret=client_secret, + purview_app_name=purview_app_name, + user_id_field=user_id_field, + guardrail_name=guardrail_name, + supported_event_hooks=supported_event_hooks, + **kwargs, + ) + self.guardrail_provider = "microsoft_purview" + verbose_proxy_logger.info( + "Initialized Microsoft Purview DLP Guardrail: %s", + guardrail_name, + ) + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + return None # Config model can be added later for UI support + + # ------------------------------------------------------------------ + # Core DLP check + # ------------------------------------------------------------------ + + async def _check_content( + self, + user_id: str, + text: str, + activity: str, + request_data: Dict[str, Any], + block_on_violation: bool = True, + ) -> Dict[str, Any]: + """Evaluate content against Purview DLP policies. + + Args: + user_id: Entra object ID. + text: Content to evaluate. + activity: ``"uploadText"`` or ``"downloadText"``. + request_data: Original request dict (used for logging metadata). + block_on_violation: If False, log only — do not raise. + + Returns: + The processContent response dict. + """ + start_time = datetime.now() + status: GuardrailStatus = "success" + response: Dict[str, Any] = {} + + try: + etag, _ = await self._compute_protection_scopes(user_id) + correlation_id = request_data.get("litellm_call_id") or str(uuid.uuid4()) + response = await self._process_content( + user_id=user_id, + text=text, + activity=activity, + etag=etag, + correlation_id=correlation_id, + ) + + if self._should_block(response): + status = "guardrail_intervened" + except HTTPException: + status = "guardrail_failed_to_respond" + raise + except httpx.HTTPStatusError as exc: + # Preserve the upstream Graph API status code (e.g. 429, 503) so + # callers can distinguish a transient infrastructure error from a + # DLP policy block (signaled separately as HTTP 400 below) and can + # implement retry-after handling on rate limits. 401/403 upstream + # responses indicate a proxy-side credential / consent problem the + # caller can do nothing about, so they are mapped to 502. + status = "guardrail_failed_to_respond" + if block_on_violation: + upstream_status = exc.response.status_code + client_status = ( + 502 if upstream_status in (401, 403) else upstream_status + ) + headers: Optional[Dict[str, str]] = None + retry_after = exc.response.headers.get("retry-after") + if retry_after: + headers = {"Retry-After": retry_after} + raise HTTPException( + status_code=client_status, + detail={ + "error": "Microsoft Purview DLP: upstream policy evaluation failed", + "activity": activity, + "upstream_status": upstream_status, + "exception": str(exc), + }, + headers=headers, + ) from exc + verbose_proxy_logger.warning( + "Purview DLP: API/network error in logging-only mode (not re-raised): %s", + exc, + ) + except Exception as exc: + status = "guardrail_failed_to_respond" + if block_on_violation: + raise HTTPException( + status_code=400, + detail={ + "error": "Microsoft Purview DLP: upstream policy evaluation failed", + "activity": activity, + "exception": str(exc), + }, + ) from exc + verbose_proxy_logger.warning( + "Purview DLP: API/network error in logging-only mode (not re-raised): %s", + exc, + ) + finally: + end_time = datetime.now() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=response, + request_data=request_data, + guardrail_status=status, + start_time=start_time.timestamp(), + end_time=end_time.timestamp(), + duration=(end_time - start_time).total_seconds(), + ) + + if block_on_violation and status == "guardrail_intervened": + raise HTTPException( + status_code=400, + detail={ + "error": "Microsoft Purview DLP: Content blocked by policy", + "activity": activity, + }, + ) + + return response + + @staticmethod + def _extract_responses_api_function_call_args(result: Any) -> List[str]: + """Return tool-call argument strings from a ``ResponsesAPIResponse.output``. + + ``ResponsesAPIResponse.output_text`` only aggregates ``output_text`` + content blocks and ignores ``function_call`` items. Model-generated + tool-call arguments can themselves contain sensitive data, so we + extract them explicitly to keep DLP coverage consistent with the + chat (``ModelResponse``) path. + """ + args: List[str] = [] + output = getattr(result, "output", None) + if not output: + return args + for item in output: + if isinstance(item, dict): + item_type = item.get("type") + arguments = item.get("arguments") + else: + item_type = getattr(item, "type", None) + arguments = getattr(item, "arguments", None) + if item_type == "function_call" and isinstance(arguments, str): + if arguments.strip(): + args.append(arguments) + return args + + def _completion_response_text_parts(self, result: Any) -> List[str]: + """Collect non-empty text segments from chat, text completions, or responses API. + + Includes assistant message content *and* model-generated tool-call + arguments so that sensitive data returned inside function calls is not + missed by the DLP scan. + """ + parts: List[str] = [] + if isinstance(result, TextCompletionResponse) and result.choices: + for text_choice in result.choices: + if not isinstance(text_choice, TextChoices): + continue + raw = text_choice.get("text") + if isinstance(raw, str) and raw.strip(): + parts.append(raw) + elif isinstance(result, ResponsesAPIResponse): + text = result.output_text + if text and text.strip(): + parts.append(text) + # Include tool-call arguments from ``function_call`` output items + # (``output_text`` ignores them). + parts.extend(self._extract_responses_api_function_call_args(result)) + elif isinstance(result, ModelResponse) and result.choices: + for chat_choice in result.choices: + if not isinstance(chat_choice, Choices): + continue + msg = chat_choice.message + if msg is None: + continue + raw = ( + msg.get("content") + if isinstance(msg, dict) + else getattr(msg, "content", None) + ) + if isinstance(raw, str) and raw.strip(): + parts.append(raw) + # Include tool-call arguments returned by the model + parts.extend(self._extract_tool_call_args_from_message(msg)) + return parts + + def _assemble_responses_api_from_chunks( + self, chunks: List[Any] + ) -> Tuple[bool, Optional[ResponsesAPIResponse]]: + """Extract the final ``ResponsesAPIResponse`` from a buffered Responses API stream. + + Returns a ``(is_responses_api_stream, assembled)`` tuple so the caller + can distinguish "not a Responses API stream" (fall through to + ``stream_chunk_builder``) from "Responses API stream but no final + response event was received" (fail closed with an accurate error). + When the stream is a Responses API stream the latest event carrying a + ``ResponsesAPIResponse`` body is returned (``response.completed``, or + ``response.failed`` / ``response.incomplete`` as fallbacks). + """ + looks_like_responses_api = False + final: Optional[ResponsesAPIResponse] = None + for chunk in chunks: + event_type = getattr(chunk, "type", None) + if isinstance(event_type, str) and event_type.startswith("response."): + looks_like_responses_api = True + candidate = getattr(chunk, "response", None) + if isinstance(candidate, ResponsesAPIResponse): + final = candidate + return looks_like_responses_api, final + + def _responses_api_input_to_str( + self, data: Dict[str, Any], raise_on_failure: bool = False + ) -> Optional[str]: + """Extract DLP-scannable text from a Responses API request ``input`` field. + + ``input`` may be a plain string or a list of input items (messages). In + the latter case the items are converted to chat messages via the standard + LiteLLM transformation and then concatenated by ``get_prompt_text_for_dlp``. + + When ``raise_on_failure`` is True (blocking mode), a transformation error + raises ``HTTPException`` so the request is fail-closed. In logging-only + mode the error is swallowed and ``None`` is returned so audit attempts on + the response side can still run. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + input_data = data.get("input") + if input_data is None and not data.get("instructions"): + return None + try: + # Always transform via messages so ``instructions`` become a system message + # (string ``input`` alone would skip instructions and bypass DLP). + messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input_data if input_data is not None else "", + responses_api_request=data, + ) + return self.get_prompt_text_for_dlp(cast(List[Any], messages)) + except Exception: + verbose_proxy_logger.warning( + "Purview DLP: failed to transform responses API input", + exc_info=True, + ) + if raise_on_failure: + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Microsoft Purview DLP: Responses API input could " + "not be transformed for DLP scanning in blocking mode" + ), + }, + ) + return None + + # ------------------------------------------------------------------ + # Identity resolution for blocking modes + # ------------------------------------------------------------------ + + def _resolve_user_id_for_blocking( + self, + data: Dict[str, Any], + user_api_key_dict: Any, + ) -> str: + """Resolve user ID for blocking (pre_call / post_call) DLP hooks. + + Uses only trusted proxy-authenticated sources (``_resolve_trusted_user_id``). + Caller-supplied ``UserAPIKeyAuth.end_user_id`` (from request ``user``, + ``metadata.user_id``, ``safety_identifier``, etc.) and + ``metadata[user_id_field]`` are rejected (fail closed) because they can + impersonate another Entra user's Purview policy. + + Raises ``HTTPException`` when no API-key-bound ``user_id`` exists or when + only caller-influenceable identity fields are available (fail closed). + """ + trusted_id = self._resolve_trusted_user_id(data, user_api_key_dict) + if trusted_id: + return trusted_id + + if self._resolve_user_id(data, user_api_key_dict): + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Microsoft Purview DLP: No proxy-authenticated user identity; " + "bind user_id to the API key (caller-supplied metadata cannot " + "be used for blocking DLP)" + ), + }, + ) + + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Microsoft Purview DLP: No proxy-authenticated user identity; " + "bind user_id to the API key for blocking DLP" + ), + }, + ) + + # ------------------------------------------------------------------ + # Pre-call hook — DLP on prompts + # ------------------------------------------------------------------ + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: Any, + data: Dict[str, Any], + call_type: "CallTypesLiteral", + ) -> Optional[Dict[str, Any]]: + """Check user prompt against Purview DLP policies before LLM call.""" + user_id = self._resolve_user_id_for_blocking(data, user_api_key_dict) + + prompt_text: Optional[str] = None + if call_type in ("responses", "aresponses"): + # Route Responses API calls to the responses-specific extractor + # before the generic ``messages`` branch. This mirrors + # ``async_logging_hook`` and ensures ``instructions`` (system + # prompt) content is included in the DLP scan, and prevents a + # crafted ``messages`` key in the request from being scanned in + # place of the actual ``input``. + prompt_text = self._responses_api_input_to_str(data, raise_on_failure=True) + elif call_type in ("text_completion", "atext_completion"): + raw_prompt = data.get("prompt") + # Reject every token-id prompt shape Purview cannot evaluate — + # flat ``list[int]`` (single prompt), ``list[list[int]]`` (multi-prompt + # batches), and mixed lists that include any token-id sub-array. + # Empty/whitespace-only strings also yield ``prompt_text is None`` but + # contain no sensitive data and pass through harmlessly below. + if self.is_token_id_prompt(raw_prompt): + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Microsoft Purview DLP: Token-id completion prompts " + "cannot be scanned for DLP in blocking mode" + ), + }, + ) + prompt_text = self.completion_prompt_to_str(raw_prompt) + else: + messages: Optional[List] = data.get("messages") + if messages: + prompt_text = self.get_prompt_text_for_dlp(cast(List[Any], messages)) + + if not prompt_text: + return data + + await self._check_content( + user_id=user_id, + text=prompt_text, + activity="uploadText", + request_data=data, + block_on_violation=True, + ) + return data + + # ------------------------------------------------------------------ + # Post-call hook — DLP on responses + # ------------------------------------------------------------------ + + @log_guardrail_information + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: "UserAPIKeyAuth", + response: Union[Any, ModelResponse, "EmbeddingResponse", "ImageResponse"], + ) -> Any: + """Check LLM response against Purview DLP policies (non-streaming only). + + Streaming responses are handled by ``async_post_call_streaming_iterator_hook`` + which buffers all chunks before scanning. The proxy automatically skips + this hook for requests that have a streaming iterator hook defined. + """ + user_id = self._resolve_user_id_for_blocking(data, user_api_key_dict) + + parts = self._completion_response_text_parts(response) + + if parts: + combined = "\n\n---\n\n".join(parts) + await self._check_content( + user_id=user_id, + text=combined, + activity="downloadText", + request_data=data, + block_on_violation=True, + ) + return response + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """Check streaming LLM responses against Purview DLP policies. + + All chunks are buffered before the DLP scan so that no content is + delivered to the client if a policy violation is detected. After a + clean scan the assembled response is re-yielded chunk-by-chunk via a + ``MockResponseIterator`` so the caller receives normal streaming output. + + The proxy automatically skips ``async_post_call_success_hook`` for + guardrails that define this method, preventing duplicate scans. + """ + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.main import stream_chunk_builder + + # Resolve user ID up-front so identity failures don't waste work + # buffering and assembling the stream. + user_id = self._resolve_user_id_for_blocking(request_data, user_api_key_dict) + + # Buffer the entire stream before any DLP scan. + all_chunks: List[ModelResponseStream] = [] + async for chunk in response: + all_chunks.append(chunk) + + # Responses API streams emit typed events (e.g. ``response.completed``) + # whose final event carries the full ``ResponsesAPIResponse`` — these + # are not understood by ``stream_chunk_builder`` (which is built for + # chat/text-completion deltas). Detect and scan them via the same + # ``_completion_response_text_parts`` path used by non-streaming. + ( + is_responses_api_stream, + responses_api_assembled, + ) = self._assemble_responses_api_from_chunks(all_chunks) + if is_responses_api_stream: + if responses_api_assembled is None: + # Fail closed: Responses API events were seen but no final + # ``response.completed`` / ``response.failed`` / + # ``response.incomplete`` event carrying a ``ResponsesAPIResponse`` + # body was received, so we cannot scan the content. + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Microsoft Purview DLP: Incomplete Responses API " + "stream — no final response event received for " + "DLP scanning; blocking response." + ), + }, + ) + parts = self._completion_response_text_parts(responses_api_assembled) + if parts: + combined = "\n\n---\n\n".join(parts) + await self._check_content( + user_id=user_id, + text=combined, + activity="downloadText", + request_data=request_data, + block_on_violation=True, + ) + for chunk in all_chunks: + yield chunk + return + + assembled_response = stream_chunk_builder(chunks=all_chunks) + + if assembled_response is None and all_chunks: + # Fail closed: stream_chunk_builder dropped all chunks, so we cannot + # scan the content. Refuse to release the buffered chunks. + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Microsoft Purview DLP: Unable to assemble streamed " + "response for scanning; blocking response." + ), + }, + ) + + if isinstance( + assembled_response, (TextCompletionResponse, ResponsesAPIResponse) + ): + parts = self._completion_response_text_parts(assembled_response) + if parts: + combined = "\n\n---\n\n".join(parts) + await self._check_content( + user_id=user_id, + text=combined, + activity="downloadText", + request_data=request_data, + block_on_violation=True, + ) + for chunk in all_chunks: + yield chunk + return + + if not isinstance(assembled_response, ModelResponse): + # Non-content response (e.g. embeddings) — pass through unchanged. + for chunk in all_chunks: + yield chunk + return + + parts = self._completion_response_text_parts(assembled_response) + if parts: + combined = "\n\n---\n\n".join(parts) + # Raises HTTPException(400) on violation — no chunks are yielded. + await self._check_content( + user_id=user_id, + text=combined, + activity="downloadText", + request_data=request_data, + block_on_violation=True, + ) + + # DLP passed — re-yield chunks from the assembled chat response. + mock_response = MockResponseIterator(model_response=assembled_response) + async for chunk in mock_response: + yield chunk + + # ------------------------------------------------------------------ + # Logging-only hook — audit without blocking + # ------------------------------------------------------------------ + + def logging_hook( + self, kwargs: dict, result: Any, call_type: str + ) -> Tuple[dict, Any]: + """Fire-and-forget async audit logging; returns original (kwargs, result) immediately. + + In the proxy's async success path, litellm independently calls both + ``logging_hook`` (sync) and ``async_logging_hook`` (async) for every + ``CustomGuardrail`` callback. To avoid making two complete sets of + Purview API calls per request, this sync hook is a no-op whenever an + event loop is running — the framework's async path will invoke + ``async_logging_hook`` directly. + + For genuine sync-only call paths (no running event loop, so the async + success handler will not fire either), schedule ``async_logging_hook`` + on a short-lived background daemon thread so audit logging still runs + without blocking the caller on two Graph API round-trips. + """ + + try: + asyncio.get_running_loop() + # Async context — let the framework's async success handler invoke + # async_logging_hook to avoid duplicate Purview API calls. Log so + # the deferral is observable if the framework ever stops dispatching + # async_logging_hook on a given code path (otherwise audit silently + # drops). + verbose_proxy_logger.debug( + "Purview audit: deferring to async_logging_hook (running event loop detected)" + ) + return kwargs, result + except RuntimeError: + pass + + async def _log_safe() -> None: + try: + await self.async_logging_hook( + kwargs=kwargs, result=result, call_type=call_type + ) + except Exception as exc: + verbose_proxy_logger.error( + "Purview audit background logging error: %s", exc + ) + + def _run_in_new_loop() -> None: + new_loop = asyncio.new_event_loop() + try: + asyncio.set_event_loop(new_loop) + new_loop.run_until_complete(_log_safe()) + finally: + new_loop.close() + asyncio.set_event_loop(None) + + thread = threading.Thread(target=_run_in_new_loop, daemon=True) + thread.start() + + return kwargs, result + + async def async_logging_hook( + self, kwargs: dict, result: Any, call_type: str + ) -> Tuple[dict, Any]: + """Send both prompt and response to Purview for audit logging. + + Errors are logged but never raised — this mode is non-blocking. + Each audit call (prompt and response) is wrapped in its own try/except + so a failure on the first does not prevent the second from running. + """ + user_id = self._resolve_user_id_from_logging_kwargs(kwargs) + if not user_id: + verbose_proxy_logger.debug("Purview audit: no user_id, skipping") + return kwargs, result + + # Log prompt (uploadText) + try: + prompt_text: Optional[str] = None + if call_type in ("responses", "aresponses"): + # Responses API: route to the responses-specific extractor + # before the generic ``messages`` branch. litellm's logging + # pipeline stores the raw responses ``input`` (a string or a + # list of input items) under ``model_call_details["messages"]`` + # via ``function_setup``, which is NOT the chat message format + # ``get_prompt_text_for_dlp`` expects. Use the original + # ``input`` / ``instructions`` keys that ``pre_call`` and + # ``update_environment_variables`` persist on the call details. + prompt_text = self._responses_api_input_to_str(kwargs) + elif call_type in ("text_completion", "atext_completion"): + prompt_text = self.completion_prompt_to_str(kwargs.get("prompt")) + else: + messages = kwargs.get("messages") + if messages: + prompt_text = self.get_prompt_text_for_dlp( + cast(List[Any], messages) + ) + + if prompt_text: + await self._check_content( + user_id=user_id, + text=prompt_text, + activity="uploadText", + request_data=kwargs, + block_on_violation=False, + ) + except Exception as e: + verbose_proxy_logger.error("Purview audit logging error (prompt): %s", e) + + # Log response (downloadText) — runs regardless of prompt audit outcome + try: + parts = self._completion_response_text_parts(result) + if parts: + combined = "\n\n---\n\n".join(parts) + await self._check_content( + user_id=user_id, + text=combined, + activity="downloadText", + request_data=kwargs, + block_on_violation=False, + ) + except Exception as e: + verbose_proxy_logger.error("Purview audit logging error (response): %s", e) + + return kwargs, result diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/base.py b/litellm/proxy/guardrails/guardrail_hooks/openai/base.py index 872d09cd886..281afacd5c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/base.py @@ -1,5 +1,9 @@ from typing import TYPE_CHECKING, List, Optional +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_last_user_message, +) + if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -21,32 +25,4 @@ def get_user_prompt(self, messages: List["AllMessageValues"]) -> Optional[str]: ] get_user_prompt(messages) -> "What is the weather in Tokyo?" """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - convert_content_list_to_str, - ) - - if not messages: - return None - - # Iterate from the end to find the last consecutive block of user messages - user_messages = [] - for message in reversed(messages): - if message.get("role") == "user": - user_messages.append(message) - else: - # Stop when we hit a non-user message - break - - if not user_messages: - return None - - # Reverse to get the messages in chronological order - user_messages.reverse() - - user_prompt = "" - for message in user_messages: - text_content = convert_content_list_to_str(message) - user_prompt += text_content + "\n" - - result = user_prompt.strip() - return result if result else None + return get_last_user_message(messages) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 0a51ce3d456..0430c570e14 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,6 +5,9 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator from typing_extensions import Required, TypedDict +from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( + AktoConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import ( BlockCodeExecutionGuardrailConfigModel, ) @@ -17,9 +20,6 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( - AktoConfigModel, -) from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) @@ -93,6 +93,7 @@ class SupportedGuardrailIntegrations(Enum): GENERIC_GUARDRAIL_API = "generic_guardrail_api" QUALIFIRE = "qualifire" CUSTOM_CODE = "custom_code" + MICROSOFT_PURVIEW = "microsoft_purview" SEMANTIC_GUARD = "semantic_guard" MCP_END_USER_PERMISSION = "mcp_end_user_permission" BLOCK_CODE_EXECUTION = "block_code_execution" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py new file mode 100644 index 00000000000..cc89cea58d2 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -0,0 +1,2659 @@ +"""Unit tests for the Microsoft Purview DLP guardrail.""" + +import asyncio +import time +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.base import ( + PurviewGuardrailBase, +) +from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import ( + MicrosoftPurviewDLPGuardrail, +) + + +def _make_guardrail(**kwargs) -> MicrosoftPurviewDLPGuardrail: + """Helper to construct a guardrail with test defaults.""" + defaults = { + "guardrail_name": "test-purview", + "tenant_id": "test-tenant-id", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + } + defaults.update(kwargs) + return MicrosoftPurviewDLPGuardrail(**defaults) + + +def _mock_token_response(): + """Mock a successful OAuth2 token response.""" + resp = Mock() + resp.json.return_value = { + "access_token": "mock-access-token", + "expires_in": 3600, + } + return resp + + +def _mock_graph_response(policy_actions=None, protection_scope_state="unchanged"): + """Mock a processContent Graph API response.""" + resp = Mock() + body = { + "protectionScopeState": protection_scope_state, + "policyActions": policy_actions or [], + "processingErrors": [], + } + resp.json.return_value = body + resp.headers = {"ETag": "test-etag-123"} + return resp + + +def _mock_scope_response(): + """Mock a protectionScopes/compute Graph API response.""" + resp = Mock() + resp.json.return_value = { + "value": [ + { + "activities": "uploadText,downloadText", + "executionMode": "evaluateInline", + "policyActions": [], + } + ] + } + resp.headers = {"ETag": "scope-etag-123"} + return resp + + +# --------------------------------------------------------------- +# _should_block +# --------------------------------------------------------------- + + +class TestShouldBlock: + def test_empty_policy_actions(self): + assert PurviewGuardrailBase._should_block({"policyActions": []}) is False + + def test_no_policy_actions_key(self): + assert PurviewGuardrailBase._should_block({}) is False + + def test_restrict_access_block(self): + response = { + "policyActions": [ + { + "@odata.type": "#microsoft.graph.restrictAccessAction", + "action": "restrictAccess", + "restrictionAction": "block", + } + ] + } + assert PurviewGuardrailBase._should_block(response) is True + + def test_restrict_access_non_block(self): + response = { + "policyActions": [ + { + "@odata.type": "#microsoft.graph.restrictAccessAction", + "action": "restrictAccess", + "restrictionAction": "warn", + } + ] + } + assert PurviewGuardrailBase._should_block(response) is False + + def test_non_restrict_action(self): + response = { + "policyActions": [ + { + "@odata.type": "#microsoft.graph.auditAction", + "action": "audit", + } + ] + } + assert PurviewGuardrailBase._should_block(response) is False + + def test_multiple_actions_one_blocks(self): + response = { + "policyActions": [ + {"action": "audit"}, + { + "@odata.type": "#microsoft.graph.restrictAccessAction", + "action": "restrictAccess", + "restrictionAction": "block", + }, + ] + } + assert PurviewGuardrailBase._should_block(response) is True + + +# --------------------------------------------------------------- +# completion prompt normalization (text completions API) +# --------------------------------------------------------------- + + +class TestCompletionPromptToStr: + def test_string_prompt(self): + assert PurviewGuardrailBase.completion_prompt_to_str(" hi ") == "hi" + + def test_list_of_strings(self): + assert PurviewGuardrailBase.completion_prompt_to_str(["a", "b"]) == "a\nb" + + def test_token_ids_returns_none(self): + assert PurviewGuardrailBase.completion_prompt_to_str([1, 2, 3]) is None + + def test_empty(self): + assert PurviewGuardrailBase.completion_prompt_to_str("") is None + assert PurviewGuardrailBase.completion_prompt_to_str([]) is None + + +# --------------------------------------------------------------- +# User ID resolution +# --------------------------------------------------------------- + + +class TestResolveUserId: + def test_from_metadata_when_no_auth_identity(self): + guardrail = _make_guardrail() + data = {"metadata": {"user_id": "entra-user-123"}} + auth = UserAPIKeyAuth(api_key="test-key-no-user") + assert guardrail._resolve_user_id(data, auth) == "entra-user-123" + + def test_authenticated_user_id_overrides_metadata(self): + """Key user_id must win over spoofed metadata[user_id_field].""" + guardrail = _make_guardrail() + data = {"metadata": {"user_id": "spoofed-entra-id"}} + auth = UserAPIKeyAuth(api_key="test", user_id="real-entra-id") + assert guardrail._resolve_user_id(data, auth) == "real-entra-id" + + def test_user_api_key_metadata_before_custom_field(self): + """Proxy-injected user_api_key_user_id wins over arbitrary metadata field.""" + guardrail = _make_guardrail(user_id_field="entra_id") + data = { + "metadata": { + "user_api_key_user_id": "from-proxy-111", + "entra_id": "metadata-222", + } + } + auth = UserAPIKeyAuth(api_key="test") + assert guardrail._resolve_user_id(data, auth) == "from-proxy-111" + + def test_custom_field_when_no_stronger_source(self): + guardrail = _make_guardrail(user_id_field="entra_id") + data = {"metadata": {"entra_id": "custom-user-456"}} + auth = UserAPIKeyAuth(api_key="test") + assert guardrail._resolve_user_id(data, auth) == "custom-user-456" + + def test_from_user_api_key_dict_user_id(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test", user_id="key-user-789") + assert guardrail._resolve_user_id({}, auth) == "key-user-789" + + def test_from_end_user_id(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test", end_user_id="end-user-101") + assert guardrail._resolve_user_id({}, auth) == "end-user-101" + + def test_end_user_id_after_key_user_id(self): + """When both key user_id and end_user_id exist, key user_id is used first.""" + guardrail = _make_guardrail() + auth = UserAPIKeyAuth( + api_key="test", user_id="key-owner", end_user_id="end-user-101" + ) + assert guardrail._resolve_user_id({}, auth) == "key-owner" + + def test_none_when_missing(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test") + assert guardrail._resolve_user_id({}, auth) is None + + +# --------------------------------------------------------------- +# Pre-call hook +# --------------------------------------------------------------- + + +class TestPreCallHook: + @pytest.mark.asyncio + async def test_pre_call_allow(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={"messages": [{"role": "user", "content": "Hello, how are you?"}]}, + call_type="completion", + ) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["activity"] == "uploadText" + assert mock_check.call_args.kwargs["block_on_violation"] is True + + @pytest.mark.asyncio + async def test_pre_call_success_returns_request_data(self): + """After a successful DLP pass, the hook must return the same data dict (not None).""" + guardrail = _make_guardrail() + payload = { + "messages": [{"role": "user", "content": "Hello, how are you?"}], + "litellm_call_id": "call-abc", + } + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + out = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data=payload, + call_type="completion", + ) + + assert out is payload + + @pytest.mark.asyncio + async def test_pre_call_block(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.side_effect = HTTPException( + status_code=400, + detail={"error": "Microsoft Purview DLP: Content blocked by policy"}, + ) + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key="test", user_id="user-123" + ), + cache=None, + data={ + "messages": [ + { + "role": "user", + "content": "SSN: 123-45-6789", + } + ] + }, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_pre_call_no_user_id_raises(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=None, + data={"messages": [{"role": "user", "content": "Hello"}]}, + call_type="completion", + ) + + mock_check.assert_not_called() + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_pre_call_no_messages_skips(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={}, + call_type="completion", + ) + + mock_check.assert_not_called() + + +class TestPreCallFullTranscript: + @pytest.mark.asyncio + async def test_pre_call_sends_all_message_roles_to_dlp(self): + """DLP text must include system / prior turns, not only the last user block.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={ + "messages": [ + {"role": "system", "content": "SYSTEM_SENSITIVE"}, + {"role": "user", "content": "EARLIER_USER"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "final benign"}, + ] + }, + call_type="completion", + ) + + mock_check.assert_called_once() + sent = mock_check.call_args.kwargs["text"] + assert "SYSTEM_SENSITIVE" in sent + assert "EARLIER_USER" in sent + assert "final benign" in sent + + +# --------------------------------------------------------------- +# Post-call hook +# --------------------------------------------------------------- + + +class TestPostCallHook: + @pytest.mark.asyncio + async def test_post_call_allow(self): + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, message=Message(content="Safe response", role="assistant") + ) + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + result = await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["activity"] == "downloadText" + assert result is response + + @pytest.mark.asyncio + async def test_post_call_block(self): + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message( + content="Credit card: 4532-6677-8521-3500", + role="assistant", + ), + ) + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.side_effect = HTTPException( + status_code=400, + detail={"error": "Microsoft Purview DLP: Content blocked by policy"}, + ) + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth( + api_key="test", user_id="user-123" + ), + response=response, + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_post_call_no_user_id_raises(self): + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices(index=0, message=Message(content="Response", role="assistant")) + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + + mock_check.assert_not_called() + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_post_call_scans_all_choices(self): + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message(content="First completion", role="assistant"), + ), + Choices( + index=1, + message=Message(content="Second completion body", role="assistant"), + ), + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + mock_check.assert_called_once() + combined = mock_check.call_args.kwargs["text"] + assert "First completion" in combined + assert "Second completion body" in combined + + +class TestTextCompletionHooks: + @pytest.mark.asyncio + async def test_pre_call_text_completion_uses_prompt(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={"prompt": "Completions API prompt body"}, + call_type="text_completion", + ) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["text"] == "Completions API prompt body" + assert mock_check.call_args.kwargs["activity"] == "uploadText" + + @pytest.mark.asyncio + async def test_post_call_text_completion_all_choices(self): + from litellm.types.utils import TextChoices, TextCompletionResponse + + guardrail = _make_guardrail() + response = TextCompletionResponse( + model="gpt-3.5-turbo-instruct", + choices=[ + TextChoices(text="alpha", index=0), + TextChoices(text="beta", index=1), + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + mock_check.assert_called_once() + combined = mock_check.call_args.kwargs["text"] + assert "alpha" in combined + assert "beta" in combined + + +# --------------------------------------------------------------- +# Responses API hooks +# --------------------------------------------------------------- + + +class TestResponsesAPIHooks: + @pytest.mark.asyncio + async def test_pre_call_responses_api_string_input(self): + """Pre-call hook must scan plain-string ``input`` on responses call type.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={"input": "SSN: 123-45-6789"}, + call_type="responses", + ) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["activity"] == "uploadText" + assert "SSN: 123-45-6789" in mock_check.call_args.kwargs["text"] + + @pytest.mark.asyncio + async def test_pre_call_aresponses_string_input(self): + """Pre-call hook must scan ``input`` on ``aresponses`` call type too.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={"input": "sensitive content"}, + call_type="aresponses", + ) + + mock_check.assert_called_once() + assert "sensitive content" in mock_check.call_args.kwargs["text"] + + @pytest.mark.asyncio + async def test_pre_call_responses_api_list_input(self): + """Pre-call hook must extract text from structured list ``input``.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={ + "input": [{"role": "user", "content": "Secret phrase: alpha bravo"}] + }, + call_type="responses", + ) + + mock_check.assert_called_once() + assert "Secret phrase: alpha bravo" in mock_check.call_args.kwargs["text"] + + @pytest.mark.asyncio + async def test_pre_call_responses_api_no_input_skips(self): + """Pre-call hook must not call _check_content when ``input`` is absent.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={}, + call_type="responses", + ) + + mock_check.assert_not_called() + + @pytest.mark.asyncio + async def test_pre_call_responses_string_input_includes_instructions(self): + """Benign string ``input`` must still scan ``instructions`` (system message).""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={ + "input": "benign user text", + "instructions": "SYSTEM_SENSITIVE in instructions", + }, + call_type="responses", + ) + + mock_check.assert_called_once() + sent = mock_check.call_args.kwargs["text"] + assert "benign user text" in sent + assert "SYSTEM_SENSITIVE in instructions" in sent + + @pytest.mark.asyncio + async def test_pre_call_responses_instructions_only(self): + """Requests with only ``instructions`` (no ``input``) must still be scanned.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={"instructions": "policy text in instructions only"}, + call_type="responses", + ) + + mock_check.assert_called_once() + assert ( + "policy text in instructions only" + in mock_check.call_args.kwargs["text"] + ) + + @pytest.mark.asyncio + async def test_post_call_responses_api_output_text(self): + """Post-call hook must scan text from ``ResponsesAPIResponse.output``.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _make_guardrail() + response = ResponsesAPIResponse( + id="resp-1", + created_at=0, + output=[ + { + "type": "message", + "id": "msg-1", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "card 4111-1111-1111-1111"} + ], + } + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + result = await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["activity"] == "downloadText" + assert "card 4111-1111-1111-1111" in mock_check.call_args.kwargs["text"] + assert result is response + + @pytest.mark.asyncio + async def test_post_call_responses_api_empty_output_skips(self): + """Post-call hook must not call _check_content when output has no text.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _make_guardrail() + response = ResponsesAPIResponse( + id="resp-2", + created_at=0, + output=[], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + mock_check.assert_not_called() + + @pytest.mark.asyncio + async def test_logging_hook_responses_api_input_and_output(self): + """Logging hook must scan both ``input`` and ``ResponsesAPIResponse.output``.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _make_guardrail() + result_response = ResponsesAPIResponse( + id="resp-3", + created_at=0, + output=[ + { + "type": "message", + "id": "msg-2", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "response body"}], + } + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_logging_hook( + kwargs={ + "input": "prompt body", + "litellm_params": { + "metadata": {"user_api_key_user_id": "user-123"} + }, + }, + result=result_response, + call_type="responses", + ) + + assert mock_check.call_count == 2 + activities = {c.kwargs["activity"] for c in mock_check.call_args_list} + assert activities == {"uploadText", "downloadText"} + texts = {c.kwargs["text"] for c in mock_check.call_args_list} + assert any("prompt body" in t for t in texts) + assert any("response body" in t for t in texts) + + @pytest.mark.asyncio + async def test_logging_hook_responses_api_with_messages_key_set(self): + """Responses-API prompt audit must fire even when ``kwargs["messages"]`` is + also set to the raw responses input. + + litellm's logging pipeline (``function_setup`` + + ``update_environment_variables``) stores the raw responses ``input`` + under ``model_call_details["messages"]``. The audit must still extract + the prompt via the responses-specific path, not silently fall through + the generic ``messages`` branch with the wrong format. + """ + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _make_guardrail() + result_response = ResponsesAPIResponse( + id="resp-msgkey", + created_at=0, + output=[ + { + "type": "message", + "id": "msg-3", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "response body"}], + } + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_logging_hook( + kwargs={ + "input": "prompt body", + "instructions": "system instructions", + # Simulate litellm's logging path which mirrors the raw + # responses input under "messages". + "messages": "prompt body", + "litellm_params": { + "metadata": {"user_api_key_user_id": "user-123"} + }, + }, + result=result_response, + call_type="aresponses", + ) + + assert mock_check.call_count == 2 + activities = {c.kwargs["activity"] for c in mock_check.call_args_list} + assert activities == {"uploadText", "downloadText"} + upload_calls = [ + c + for c in mock_check.call_args_list + if c.kwargs["activity"] == "uploadText" + ] + assert len(upload_calls) == 1 + upload_text = upload_calls[0].kwargs["text"] + assert "prompt body" in upload_text + assert "system instructions" in upload_text + + +# --------------------------------------------------------------- +# Logging hook user resolution +# --------------------------------------------------------------- + + +class TestLoggingResolveUserId: + def test_logging_prefers_user_api_key_user_id_in_metadata(self): + guardrail = _make_guardrail() + kwargs = { + "litellm_params": { + "metadata": { + "user_api_key_user_id": "trusted-from-proxy", + "user_id": "metadata-spoof", + } + } + } + assert ( + guardrail._resolve_user_id_from_logging_kwargs(kwargs) + == "trusted-from-proxy" + ) + + def test_logging_ignores_caller_supplied_user_id_field(self): + """Caller-controlled ``metadata[user_id_field]`` must not drive Purview audit attribution.""" + guardrail = _make_guardrail() + kwargs = {"litellm_params": {"metadata": {"user_id": "only-metadata-user"}}} + assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) is None + + def test_logging_kwargs_level_user_api_key_user_id(self): + """Top-level ``kwargs["user_api_key_user_id"]`` is also a proxy-injected source.""" + guardrail = _make_guardrail() + kwargs = { + "user_api_key_user_id": "from-top-level", + "litellm_params": {"metadata": {}}, + } + assert ( + guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "from-top-level" + ) + + def test_logging_returns_none_when_no_trusted_identity(self): + guardrail = _make_guardrail() + kwargs = {"litellm_params": {"metadata": {}}} + assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) is None + + +# --------------------------------------------------------------- +# _check_content — integration-level +# --------------------------------------------------------------- + + +class TestCheckContent: + @pytest.mark.asyncio + async def test_check_content_allow(self): + guardrail = _make_guardrail() + + with ( + patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + return_value=("etag-1", {}), + ), + patch.object( + guardrail, + "_process_content", + new_callable=AsyncMock, + return_value={ + "protectionScopeState": "unchanged", + "policyActions": [], + }, + ), + ): + result = await guardrail._check_content( + user_id="user-1", + text="Hello world", + activity="uploadText", + request_data={}, + block_on_violation=True, + ) + + assert result["policyActions"] == [] + + @pytest.mark.asyncio + async def test_check_content_block(self): + guardrail = _make_guardrail() + + with ( + patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + return_value=("etag-1", {}), + ), + patch.object( + guardrail, + "_process_content", + new_callable=AsyncMock, + return_value={ + "protectionScopeState": "unchanged", + "policyActions": [ + { + "@odata.type": "#microsoft.graph.restrictAccessAction", + "action": "restrictAccess", + "restrictionAction": "block", + } + ], + }, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail._check_content( + user_id="user-1", + text="SSN: 123-45-6789", + activity="uploadText", + request_data={}, + block_on_violation=True, + ) + + assert exc_info.value.status_code == 400 + assert "blocked by policy" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_check_content_logging_only_no_block(self): + """In logging_only mode, violations should NOT raise.""" + guardrail = _make_guardrail() + + with ( + patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + return_value=("etag-1", {}), + ), + patch.object( + guardrail, + "_process_content", + new_callable=AsyncMock, + return_value={ + "protectionScopeState": "unchanged", + "policyActions": [ + { + "@odata.type": "#microsoft.graph.restrictAccessAction", + "action": "restrictAccess", + "restrictionAction": "block", + } + ], + }, + ), + ): + # Should NOT raise even though violation detected + result = await guardrail._check_content( + user_id="user-1", + text="SSN: 123-45-6789", + activity="uploadText", + request_data={}, + block_on_violation=False, + ) + + assert len(result["policyActions"]) == 1 + + +# --------------------------------------------------------------- +# Token caching +# --------------------------------------------------------------- + + +class TestTokenCaching: + @pytest.mark.asyncio + async def test_token_cached(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail.async_handler, "post", return_value=_mock_token_response() + ) as mock_post: + token1 = await guardrail._get_access_token() + token2 = await guardrail._get_access_token() + + assert token1 == "mock-access-token" + assert token2 == "mock-access-token" + # Should only call the token endpoint once (cached) + assert mock_post.call_count == 1 + + @pytest.mark.asyncio + async def test_token_refreshed_on_expiry(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail.async_handler, "post", return_value=_mock_token_response() + ) as mock_post: + await guardrail._get_access_token() + + # Expire the token + guardrail._token_cache = ("old-token", time.time() - 10) + + await guardrail._get_access_token() + + # Should have called token endpoint twice + assert mock_post.call_count == 2 + + @pytest.mark.asyncio + async def test_token_http_error_propagates(self): + """OAuth2 4xx/5xx responses must surface as HTTPStatusError, not KeyError.""" + guardrail = _make_guardrail() + + error_resp = Mock() + error_resp.json.return_value = { + "error": "invalid_client", + "error_description": "client secret is wrong", + } + error_resp.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "401 Unauthorized", + request=httpx.Request("POST", "https://login.microsoftonline.com/"), + response=httpx.Response(401), + ) + ) + + with patch.object(guardrail.async_handler, "post", return_value=error_resp): + with pytest.raises(httpx.HTTPStatusError): + await guardrail._get_access_token() + + # Failure must not poison the cache. + assert guardrail._token_cache is None + + +# --------------------------------------------------------------- +# Graph POST HTTP error propagation +# --------------------------------------------------------------- + + +class TestGraphPostHttpError: + @pytest.mark.asyncio + async def test_graph_post_http_error_propagates(self): + """Non-2xx Graph API responses must raise rather than return error JSON.""" + guardrail = _make_guardrail() + + error_resp = Mock() + error_resp.json.return_value = { + "error": {"code": "Forbidden", "message": "no access"} + } + error_resp.headers = {} + error_resp.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "403 Forbidden", + request=httpx.Request("POST", "https://graph.microsoft.com/"), + response=httpx.Response(403), + ) + ) + + with ( + patch.object( + guardrail, "_get_access_token", new_callable=AsyncMock + ) as mock_token, + patch.object(guardrail.async_handler, "post", return_value=error_resp), + ): + mock_token.return_value = "mock-token" + + with pytest.raises(httpx.HTTPStatusError): + await guardrail._graph_post( + "https://graph.microsoft.com/v1.0/users/u/example", + {"foo": "bar"}, + ) + + @pytest.mark.asyncio + async def test_compute_protection_scopes_http_error_propagates(self): + """A Graph error on protectionScopes/compute must not be cached as success.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_graph_post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = httpx.HTTPStatusError( + "429 Too Many Requests", + request=httpx.Request("POST", "https://graph.microsoft.com/"), + response=httpx.Response(429), + ) + + with pytest.raises(httpx.HTTPStatusError): + await guardrail._compute_protection_scopes("user-err") + + # The failed compute must not populate the scope cache. + assert "user-err" not in guardrail._scope_cache + + +# --------------------------------------------------------------- +# Protection scope caching +# --------------------------------------------------------------- + + +class TestScopeCaching: + @pytest.mark.asyncio + async def test_scope_cached(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_graph_post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = ( + { + "value": [ + {"activities": "uploadText", "executionMode": "evaluateInline"} + ] + }, + {"ETag": "scope-etag"}, + ) + + etag1, _ = await guardrail._compute_protection_scopes("user-1") + etag2, _ = await guardrail._compute_protection_scopes("user-1") + + assert etag1 == "scope-etag" + assert etag2 == "scope-etag" + assert mock_post.call_count == 1 + + @pytest.mark.asyncio + async def test_scope_cache_lru_keeps_hot_user_on_eviction(self): + """Frequently accessed users should not be evicted before cold entries.""" + guardrail = _make_guardrail() + guardrail._scope_cache_maxsize = 3 + + scope_payload = ( + { + "value": [ + {"activities": "uploadText", "executionMode": "evaluateInline"} + ] + }, + {"ETag": "scope-etag"}, + ) + + with patch.object( + guardrail, "_graph_post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = scope_payload + + await guardrail._compute_protection_scopes("user-a") + await guardrail._compute_protection_scopes("user-b") + await guardrail._compute_protection_scopes("user-c") + assert mock_post.call_count == 3 + + await guardrail._compute_protection_scopes("user-a") + assert mock_post.call_count == 3 + + await guardrail._compute_protection_scopes("user-d") + assert mock_post.call_count == 4 + + await guardrail._compute_protection_scopes("user-a") + assert mock_post.call_count == 4 + assert "user-a" in guardrail._scope_cache + assert "user-b" not in guardrail._scope_cache + + @pytest.mark.asyncio + async def test_scope_cache_refresh_moves_to_end_of_lru(self): + """Refreshing a stale entry must move it to the MRU end of the OrderedDict. + + Before the fix, OrderedDict.__setitem__ preserved the original insertion + position for existing keys, causing the just-refreshed entry to be the + next candidate for LRU eviction. + """ + guardrail = _make_guardrail() + guardrail._scope_cache_maxsize = 2 + + scope_payload = ( + {"value": []}, + {"ETag": "scope-etag"}, + ) + + with patch.object( + guardrail, "_graph_post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = scope_payload + + # Populate cache: user-a (older), user-b (newer) + await guardrail._compute_protection_scopes("user-a") + await guardrail._compute_protection_scopes("user-b") + assert mock_post.call_count == 2 + + # Expire user-a's entry so it is re-fetched on the next access. + old_etag, old_scope, _ = guardrail._scope_cache["user-a"] + guardrail._scope_cache["user-a"] = (old_etag, old_scope, 0.0) + + # Re-fetch user-a — should move it to the MRU end. + await guardrail._compute_protection_scopes("user-a") + assert mock_post.call_count == 3 + + # Adding a third user must evict user-b (the true LRU), not user-a. + await guardrail._compute_protection_scopes("user-c") + assert mock_post.call_count == 4 + + assert "user-a" in guardrail._scope_cache, "user-a was wrongly evicted" + assert ( + "user-b" not in guardrail._scope_cache + ), "user-b should have been evicted" + assert "user-c" in guardrail._scope_cache + + @pytest.mark.asyncio + async def test_scope_invalidated_on_modified(self): + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_graph_post", new_callable=AsyncMock + ) as mock_post: + # First call: compute scopes + mock_post.return_value = ( + {"value": []}, + {"ETag": "etag-1"}, + ) + await guardrail._compute_protection_scopes("user-1") + + # processContent returns modified + mock_post.return_value = ( + {"protectionScopeState": "modified", "policyActions": []}, + {}, + ) + await guardrail._process_content("user-1", "text", "uploadText", "etag-1") + + # Scope cache should be invalidated + assert "user-1" not in guardrail._scope_cache + + +# --------------------------------------------------------------- +# get_prompt_text_for_dlp — message separator +# --------------------------------------------------------------- + + +class TestGetPromptTextForDlp: + def test_single_message_no_extra_separator(self): + """A single message is returned as-is (no leading/trailing separator).""" + guardrail = _make_guardrail() + result = guardrail.get_prompt_text_for_dlp( + [{"role": "user", "content": "Hello"}] + ) + assert result == "Hello" + + def test_messages_separated_by_double_newline(self): + """Adjacent messages must NOT be concatenated without a separator. + + Before the fix, "end of msg1" + "start of msg2" became + "end of msg1start of msg2", mangling DLP pattern detection. + """ + guardrail = _make_guardrail() + result = guardrail.get_prompt_text_for_dlp( + [ + {"role": "system", "content": "end of msg1"}, + {"role": "user", "content": "start of msg2"}, + ] + ) + assert result is not None + assert "end of msg1" in result + assert "start of msg2" in result + # Separator must be present between messages + assert "end of msg1start of msg2" not in result + assert "end of msg1\n\nstart of msg2" in result + + def test_empty_messages_returns_none(self): + guardrail = _make_guardrail() + assert guardrail.get_prompt_text_for_dlp([]) is None + + def test_whitespace_only_messages_skipped(self): + guardrail = _make_guardrail() + result = guardrail.get_prompt_text_for_dlp( + [ + {"role": "system", "content": " "}, + {"role": "user", "content": "real content"}, + ] + ) + assert result == "real content" + + def test_multi_role_conversation_preserves_all_content(self): + guardrail = _make_guardrail() + result = guardrail.get_prompt_text_for_dlp( + [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "USER1"}, + {"role": "assistant", "content": "ASSISTANT"}, + {"role": "user", "content": "USER2"}, + ] + ) + assert result is not None + for token in ("SYSTEM", "USER1", "ASSISTANT", "USER2"): + assert token in result + + +# --------------------------------------------------------------- +# logging_hook — non-blocking fire-and-forget +# --------------------------------------------------------------- + + +class TestLoggingHookNonBlocking: + @pytest.mark.asyncio + async def test_logging_hook_does_not_block_running_loop(self): + """logging_hook must return immediately without blocking the event loop. + + Before the fix, logging_hook called future.result() which blocked the + event loop thread for the full round-trip of the two Graph API calls. + """ + guardrail = _make_guardrail() + call_count = 0 + + async def slow_async_hook(**_kwargs): + nonlocal call_count + await asyncio.sleep(0.05) + call_count += 1 + return _kwargs.get("kwargs", {}), _kwargs.get("result") + + with patch.object(guardrail, "async_logging_hook", side_effect=slow_async_hook): + # Call logging_hook from within a running event loop + result = guardrail.logging_hook( + kwargs={"messages": [{"role": "user", "content": "test"}]}, + result=None, + call_type="completion", + ) + + # Must return (kwargs, result) unchanged without waiting for async work + assert result[0]["messages"][0]["content"] == "test" + assert result[1] is None + + def test_logging_hook_returns_original_kwargs_and_result(self): + """Return value must be the original (kwargs, result) tuple unchanged.""" + guardrail = _make_guardrail() + kwargs = {"messages": [{"role": "user", "content": "hello"}]} + result_obj = {"some": "result"} + + with patch.object( + guardrail, + "async_logging_hook", + new_callable=AsyncMock, + return_value=(kwargs, result_obj), + ): + out = guardrail.logging_hook( + kwargs=kwargs, + result=result_obj, + call_type="completion", + ) + + assert out == (kwargs, result_obj) + + +# --------------------------------------------------------------- +# Initializer validation +# --------------------------------------------------------------- + + +class TestInitializerValidation: + def test_missing_tenant_id(self): + from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview import ( + initialize_guardrail, + ) + + litellm_params = Mock( + spec=[ + "tenant_id", + "client_id", + "client_secret", + "purview_app_name", + "user_id_field", + "api_key", + "mode", + "default_on", + ] + ) + litellm_params.tenant_id = None + litellm_params.client_id = None + litellm_params.client_secret = None + litellm_params.api_key = "secret" + litellm_params.mode = "pre_call" + + with pytest.raises(ValueError, match="tenant_id is required"): + initialize_guardrail(litellm_params, {"guardrail_name": "test"}) + + def test_missing_client_id(self): + from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview import ( + initialize_guardrail, + ) + + litellm_params = Mock( + spec=[ + "tenant_id", + "client_id", + "client_secret", + "purview_app_name", + "user_id_field", + "api_key", + "mode", + "default_on", + ] + ) + litellm_params.tenant_id = "test-tenant" + litellm_params.client_id = None + litellm_params.client_secret = None + litellm_params.api_key = "secret" + litellm_params.mode = "pre_call" + + with pytest.raises(ValueError, match="client_id is required"): + initialize_guardrail(litellm_params, {"guardrail_name": "test"}) + + def test_missing_client_secret(self): + from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview import ( + initialize_guardrail, + ) + + litellm_params = Mock( + spec=[ + "tenant_id", + "client_id", + "client_secret", + "purview_app_name", + "user_id_field", + "api_key", + "mode", + "default_on", + ] + ) + litellm_params.tenant_id = "test-tenant" + litellm_params.client_id = "test-client" + litellm_params.client_secret = None + litellm_params.api_key = None + litellm_params.mode = "pre_call" + + with pytest.raises(ValueError, match="client_secret"): + initialize_guardrail(litellm_params, {"guardrail_name": "test"}) + + +# --------------------------------------------------------------- +# _check_content — API error handling with block_on_violation=False +# --------------------------------------------------------------- + + +class TestCheckContentApiErrorHandling: + @pytest.mark.asyncio + async def test_api_error_reraises_when_block_on_violation_true(self): + """API/network errors must surface as HTTPException(400) when block_on_violation=True.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + side_effect=RuntimeError("network failure"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail._check_content( + user_id="user-1", + text="some content", + activity="uploadText", + request_data={}, + block_on_violation=True, + ) + + assert exc_info.value.status_code == 400 + assert isinstance(exc_info.value.detail, dict) + assert "upstream policy evaluation failed" in exc_info.value.detail.get( + "error", "" + ) + assert "network failure" in exc_info.value.detail.get("exception", "") + assert isinstance(exc_info.value.__cause__, RuntimeError) + + @pytest.mark.asyncio + async def test_http_exception_passes_through_unchanged(self): + """HTTPException from upstream layers must propagate as-is (not wrapped).""" + guardrail = _make_guardrail() + inner = HTTPException(status_code=403, detail="forbidden") + + with patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + side_effect=inner, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail._check_content( + user_id="user-1", + text="some content", + activity="uploadText", + request_data={}, + block_on_violation=True, + ) + + assert exc_info.value is inner + + @pytest.mark.asyncio + async def test_api_error_not_reraised_when_block_on_violation_false(self): + """API/network errors must be swallowed (logged only) when block_on_violation=False.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + side_effect=RuntimeError("network failure"), + ): + # Must NOT raise — should return empty dict + result = await guardrail._check_content( + user_id="user-1", + text="some content", + activity="uploadText", + request_data={}, + block_on_violation=False, + ) + + assert isinstance(result, dict) + + @pytest.mark.asyncio + async def test_process_content_error_not_reraised_when_block_on_violation_false( + self, + ): + """Errors from _process_content itself must also be suppressed in logging-only mode.""" + guardrail = _make_guardrail() + + with ( + patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + return_value=("etag-1", {}), + ), + patch.object( + guardrail, + "_process_content", + new_callable=AsyncMock, + side_effect=ConnectionError("timeout"), + ), + ): + result = await guardrail._check_content( + user_id="user-1", + text="some content", + activity="uploadText", + request_data={}, + block_on_violation=False, + ) + + assert isinstance(result, dict) + + @pytest.mark.asyncio + async def test_http_status_error_preserves_upstream_status_code(self): + """Upstream Graph 429 must surface as 429 with Retry-After (not a generic 400).""" + guardrail = _make_guardrail() + upstream_resp = httpx.Response( + status_code=429, + headers={"Retry-After": "30"}, + request=httpx.Request("POST", "https://graph.microsoft.com/v1.0/x"), + ) + upstream_err = httpx.HTTPStatusError( + "rate limited", request=upstream_resp.request, response=upstream_resp + ) + + with patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + side_effect=upstream_err, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail._check_content( + user_id="user-1", + text="some content", + activity="uploadText", + request_data={}, + block_on_violation=True, + ) + + assert exc_info.value.status_code == 429 + assert exc_info.value.headers == {"Retry-After": "30"} + assert isinstance(exc_info.value.detail, dict) + assert exc_info.value.detail.get("upstream_status") == 429 + assert isinstance(exc_info.value.__cause__, httpx.HTTPStatusError) + + @pytest.mark.asyncio + async def test_http_status_error_401_maps_to_502(self): + """Upstream 401/403 (proxy creds problem) should be exposed as 502, not 401/403.""" + guardrail = _make_guardrail() + upstream_resp = httpx.Response( + status_code=401, + request=httpx.Request("POST", "https://graph.microsoft.com/v1.0/x"), + ) + upstream_err = httpx.HTTPStatusError( + "unauthorized", request=upstream_resp.request, response=upstream_resp + ) + + with patch.object( + guardrail, + "_compute_protection_scopes", + new_callable=AsyncMock, + side_effect=upstream_err, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail._check_content( + user_id="user-1", + text="some content", + activity="uploadText", + request_data={}, + block_on_violation=True, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail.get("upstream_status") == 401 + + +# --------------------------------------------------------------- +# async_logging_hook — independent prompt/response audit calls +# --------------------------------------------------------------- + + +class TestAsyncLoggingHookIndependence: + @pytest.mark.asyncio + async def test_response_audit_runs_even_if_prompt_audit_fails(self): + """A failure in the prompt audit must not prevent the response audit from running.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message(content="response text", role="assistant"), + ) + ], + ) + + call_activities: list = [] + + async def fake_check_content(**kwargs): + activity = kwargs.get("activity") + if activity == "uploadText": + raise RuntimeError("simulated prompt API failure") + call_activities.append(activity) + return {"policyActions": []} + + with patch.object(guardrail, "_check_content", side_effect=fake_check_content): + await guardrail.async_logging_hook( + kwargs={ + "messages": [{"role": "user", "content": "prompt"}], + "litellm_params": { + "metadata": {"user_api_key_user_id": "user-123"} + }, + }, + result=response, + call_type="completion", + ) + + # The response audit must still have been attempted + assert "downloadText" in call_activities + + @pytest.mark.asyncio + async def test_prompt_audit_runs_even_if_response_audit_fails(self): + """A failure in the response audit must not affect the prompt audit result.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message(content="response text", role="assistant"), + ) + ], + ) + + call_activities: list = [] + + async def fake_check_content(**kwargs): + activity = kwargs.get("activity") + if activity == "downloadText": + raise RuntimeError("simulated response API failure") + call_activities.append(activity) + return {"policyActions": []} + + with patch.object(guardrail, "_check_content", side_effect=fake_check_content): + await guardrail.async_logging_hook( + kwargs={ + "messages": [{"role": "user", "content": "prompt"}], + "litellm_params": { + "metadata": {"user_api_key_user_id": "user-123"} + }, + }, + result=response, + call_type="completion", + ) + + assert "uploadText" in call_activities + + @pytest.mark.asyncio + async def test_logging_hook_returns_original_when_both_audits_fail(self): + """async_logging_hook must always return (kwargs, result) even if both audits fail.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, + "_check_content", + new_callable=AsyncMock, + side_effect=RuntimeError("total failure"), + ): + kwargs = { + "messages": [{"role": "user", "content": "prompt"}], + "litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}}, + } + result_obj = {"some": "result"} + out_kwargs, out_result = await guardrail.async_logging_hook( + kwargs=kwargs, + result=result_obj, + call_type="completion", + ) + + assert out_kwargs is kwargs + assert out_result is result_obj + + +# --------------------------------------------------------------- +# Tool-call argument extraction +# --------------------------------------------------------------- + + +class TestExtractToolCallArgs: + def test_dict_message_with_tool_calls(self): + msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + {"function": {"arguments": '{"ssn": "123-45-6789"}'}}, + {"function": {"arguments": '{"card": "4111-1111-1111-1111"}'}}, + ], + } + args = MicrosoftPurviewDLPGuardrail._extract_tool_call_args_from_message(msg) + assert '{"ssn": "123-45-6789"}' in args + assert '{"card": "4111-1111-1111-1111"}' in args + + def test_dict_message_with_function_call(self): + msg = { + "role": "assistant", + "content": None, + "function_call": {"name": "lookup", "arguments": '{"query": "secret"}'}, + } + args = MicrosoftPurviewDLPGuardrail._extract_tool_call_args_from_message(msg) + assert '{"query": "secret"}' in args + + def test_object_message_with_tool_calls(self): + from litellm.types.utils import Message + + msg = Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "tc1", + "type": "function", + "function": {"name": "fn", "arguments": '{"x": 1}'}, + }, + ], + ) + args = MicrosoftPurviewDLPGuardrail._extract_tool_call_args_from_message(msg) + assert '{"x": 1}' in args + + def test_message_with_no_tool_calls(self): + msg = {"role": "user", "content": "hello"} + args = MicrosoftPurviewDLPGuardrail._extract_tool_call_args_from_message(msg) + assert args == [] + + def test_empty_arguments_skipped(self): + msg = { + "role": "assistant", + "content": None, + "tool_calls": [{"function": {"arguments": " "}}], + } + args = MicrosoftPurviewDLPGuardrail._extract_tool_call_args_from_message(msg) + assert args == [] + + +# --------------------------------------------------------------- +# Tool-call arguments included in DLP text extraction (prompt) +# --------------------------------------------------------------- + + +class TestGetPromptTextToolCalls: + def test_tool_call_args_included_in_prompt_scan(self): + """Sensitive data in tool_calls[].function.arguments must appear in DLP text.""" + guardrail = _make_guardrail() + messages = [ + {"role": "user", "content": "benign query"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tc1", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"ssn": "123-45-6789"}', + }, + } + ], + }, + ] + text = guardrail.get_prompt_text_for_dlp(messages) + assert text is not None + assert "benign query" in text + assert '{"ssn": "123-45-6789"}' in text + + def test_function_call_args_included_in_prompt_scan(self): + """Legacy function_call.arguments must also appear in DLP text.""" + guardrail = _make_guardrail() + messages = [ + { + "role": "assistant", + "content": "Calling function", + "function_call": { + "name": "search", + "arguments": '{"credit_card": "4111-1111-1111-1111"}', + }, + } + ] + text = guardrail.get_prompt_text_for_dlp(messages) + assert text is not None + assert "Calling function" in text + assert '{"credit_card": "4111-1111-1111-1111"}' in text + + def test_content_only_message_unchanged(self): + """Messages without tool calls must still produce the same output.""" + guardrail = _make_guardrail() + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Tell me a joke."}, + ] + text = guardrail.get_prompt_text_for_dlp(messages) + assert text is not None + assert "You are helpful." in text + assert "Tell me a joke." in text + + @pytest.mark.asyncio + async def test_pre_call_hook_scans_tool_call_args(self): + """async_pre_call_hook must include tool_call arguments in the text sent to Purview.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_pre_call_hook( + user_api_key_dict=__import__( + "litellm.proxy._types", fromlist=["UserAPIKeyAuth"] + ).UserAPIKeyAuth(api_key="test", user_id="user-123"), + cache=None, + data={ + "messages": [ + {"role": "user", "content": "benign"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tc1", + "type": "function", + "function": { + "name": "do_thing", + "arguments": '{"password": "hunter2"}', + }, + } + ], + }, + ] + }, + call_type="completion", + ) + + mock_check.assert_called_once() + sent_text = mock_check.call_args.kwargs["text"] + assert '{"password": "hunter2"}' in sent_text + + +# --------------------------------------------------------------- +# Tool-call arguments included in DLP text extraction (response) +# --------------------------------------------------------------- + + +class TestCompletionResponseTextPartsToolCalls: + def test_response_tool_call_args_included(self): + """Model-generated tool_call arguments must appear in the DLP scan text.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "tc1", + "type": "function", + "function": { + "name": "exfil", + "arguments": '{"data": "secret-value"}', + }, + } + ], + ), + ) + ], + ) + parts = guardrail._completion_response_text_parts(response) + assert any("secret-value" in p for p in parts) + + def test_response_with_content_and_tool_calls(self): + """Both message content and tool_call arguments must be included.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message( + role="assistant", + content="Here is the result", + tool_calls=[ + { + "id": "tc2", + "type": "function", + "function": { + "name": "fn", + "arguments": '{"ssn": "123-45-6789"}', + }, + } + ], + ), + ) + ], + ) + parts = guardrail._completion_response_text_parts(response) + combined = " ".join(parts) + assert "Here is the result" in combined + assert '{"ssn": "123-45-6789"}' in combined + + @pytest.mark.asyncio + async def test_post_call_hook_scans_response_tool_call_args(self): + """async_post_call_success_hook must send tool_call arguments to Purview.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "tc3", + "type": "function", + "function": { + "name": "retrieve", + "arguments": '{"credit_card": "4111-1111-1111-1111"}', + }, + } + ], + ), + ) + ], + ) + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + + await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=__import__( + "litellm.proxy._types", fromlist=["UserAPIKeyAuth"] + ).UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=response, + ) + + mock_check.assert_called_once() + sent_text = mock_check.call_args.kwargs["text"] + assert '{"credit_card": "4111-1111-1111-1111"}' in sent_text + + def test_responses_api_function_call_args_included(self): + """Function-call arguments in ``ResponsesAPIResponse.output`` must be DLP-scanned.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _make_guardrail() + response = ResponsesAPIResponse( + id="resp-tc-1", + created_at=0, + output=[ + { + "type": "message", + "id": "msg-tc-1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "calling tool"}], + }, + { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "exfil", + "arguments": '{"ssn": "123-45-6789"}', + }, + ], + ) + parts = guardrail._completion_response_text_parts(response) + combined = " ".join(parts) + assert "calling tool" in combined + assert '{"ssn": "123-45-6789"}' in combined + + def test_responses_api_function_call_args_only(self): + """Function-call args must be scanned even when no ``output_text`` blocks exist.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _make_guardrail() + response = ResponsesAPIResponse( + id="resp-tc-2", + created_at=0, + output=[ + { + "type": "function_call", + "id": "fc-2", + "call_id": "call-2", + "name": "exfil", + "arguments": '{"secret": "hunter2"}', + } + ], + ) + parts = guardrail._completion_response_text_parts(response) + assert any('{"secret": "hunter2"}' in p for p in parts) + + +# --------------------------------------------------------------- +# Graph user id path encoding +# --------------------------------------------------------------- + + +class TestGraphUserIdEncoding: + def test_encode_graph_user_id_percent_encodes_special_chars(self): + from urllib.parse import quote + + raw = "user/with%special" + encoded = PurviewGuardrailBase._encode_graph_user_id(raw) + assert encoded == quote(raw, safe="") + assert "/" not in encoded + + @pytest.mark.asyncio + async def test_compute_protection_scopes_uses_encoded_path(self): + guardrail = _make_guardrail() + guardrail._scope_cache.clear() + + mock_resp = _mock_scope_response() + + async def _capture_post(url, **kwargs): + assert "/users/" in url + assert "user%2Fwith%25special" in url + return mock_resp + + guardrail.async_handler.post = AsyncMock(side_effect=_capture_post) + + with patch.object( + guardrail, "_get_access_token", new_callable=AsyncMock + ) as mock_token: + mock_token.return_value = "tok" + await guardrail._compute_protection_scopes("user/with%special") + + guardrail.async_handler.post.assert_called_once() + + +# --------------------------------------------------------------- +# _resolve_trusted_user_id +# --------------------------------------------------------------- + + +class TestResolveTrustedUserId: + def test_trusted_user_id_from_api_key_dict(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test", user_id="auth-user-111") + assert guardrail._resolve_trusted_user_id({}, auth) == "auth-user-111" + + def test_end_user_id_not_trusted_for_blocking(self): + """end_user_id is request-derived; must not be used for blocking DLP.""" + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test", end_user_id="end-user-222") + assert guardrail._resolve_trusted_user_id({}, auth) is None + + def test_metadata_user_api_key_user_id_not_trusted_without_auth(self): + """Metadata user_api_key_user_id is not trusted when the key has no user_id.""" + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test") + data = {"metadata": {"user_api_key_user_id": "proxy-user-333"}} + assert guardrail._resolve_trusted_user_id(data, auth) is None + + def test_trusted_user_id_returns_none_for_caller_supplied_only(self): + """Caller-supplied metadata must NOT be returned by _resolve_trusted_user_id.""" + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test") + data = {"metadata": {"user_id": "caller-supplied-444"}} + assert guardrail._resolve_trusted_user_id(data, auth) is None + + def test_trusted_prefers_key_user_id_over_end_user_id(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth( + api_key="test", user_id="key-owner", end_user_id="end-user" + ) + assert guardrail._resolve_trusted_user_id({}, auth) == "key-owner" + + +# --------------------------------------------------------------- +# _resolve_user_id_from_logging_kwargs — caller-influenceable identity rejected +# --------------------------------------------------------------- + + +class TestLoggingRejectsCallerInfluenceableIdentity: + """``end_user_id`` is derived from caller-controllable request fields + (``user``, ``metadata.user_id``, ``safety_identifier``) so it must not + drive Purview audit attribution either. + """ + + def test_end_user_id_in_metadata_is_ignored(self): + guardrail = _make_guardrail() + kwargs = { + "litellm_params": { + "metadata": { + "user_api_key_end_user_id": "end-user-from-metadata", + } + } + } + assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) is None + + def test_end_user_id_at_top_level_kwargs_is_ignored(self): + guardrail = _make_guardrail() + kwargs = { + "user_api_key_end_user_id": "end-user-from-kwargs", + "litellm_params": {"metadata": {}}, + } + assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) is None + + +# --------------------------------------------------------------- +# _resolve_user_id_for_blocking — security warning path +# --------------------------------------------------------------- + + +class TestResolveUserIdForBlocking: + def test_trusted_id_returned_without_warning(self, caplog): + import logging + + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test", user_id="trusted-111") + with caplog.at_level(logging.WARNING): + result = guardrail._resolve_user_id_for_blocking({}, auth) + assert result == "trusted-111" + assert "SECURITY" not in caplog.text + + def test_caller_supplied_id_raises_http_exception(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test") + data = {"metadata": {"user_id": "caller-supplied-999"}} + with pytest.raises(HTTPException) as exc_info: + guardrail._resolve_user_id_for_blocking(data, auth) + assert exc_info.value.status_code == 400 + assert "proxy-authenticated" in str(exc_info.value.detail) + + def test_no_id_raises_http_exception(self): + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test") + with pytest.raises(HTTPException) as exc_info: + guardrail._resolve_user_id_for_blocking({}, auth) + assert exc_info.value.status_code == 400 + assert "bind user_id" in str(exc_info.value.detail) + + def test_end_user_id_only_raises_for_blocking(self): + """Request-derived end_user_id cannot drive blocking Purview checks.""" + guardrail = _make_guardrail() + auth = UserAPIKeyAuth(api_key="test", end_user_id="caller-end-user") + with pytest.raises(HTTPException) as exc_info: + guardrail._resolve_user_id_for_blocking({}, auth) + assert exc_info.value.status_code == 400 + assert "proxy-authenticated" in str(exc_info.value.detail) + + +# --------------------------------------------------------------- +# Token-id prompt handling in pre_call blocking mode +# --------------------------------------------------------------- + + +class TestTokenIdPromptHandling: + @pytest.mark.asyncio + async def test_token_id_prompt_raises_in_blocking_mode(self): + """Pure token-id prompts must be rejected in blocking pre_call mode.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), + cache=None, + data={"prompt": [1, 2, 3, 100, 200]}, + call_type="text_completion", + ) + + mock_check.assert_not_called() + assert exc_info.value.status_code == 400 + assert "Token-id" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_missing_prompt_skips_without_warning(self, caplog): + """No prompt at all → silently skip (not a token-id bypass case).""" + import logging + + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + with caplog.at_level(logging.WARNING): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), + cache=None, + data={}, + call_type="text_completion", + ) + + mock_check.assert_not_called() + assert "token-id" not in caplog.text.lower() + + @pytest.mark.asyncio + async def test_string_prompt_still_scanned(self): + """Normal string prompts must still be sent to Purview.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + mock_check.return_value = {"policyActions": []} + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), + cache=None, + data={"prompt": "sensitive text"}, + call_type="text_completion", + ) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["text"] == "sensitive text" + + @pytest.mark.asyncio + @pytest.mark.parametrize("empty_prompt", ["", " ", "\n\t "]) + async def test_empty_or_whitespace_prompt_passes_through(self, empty_prompt): + """Empty/whitespace-only string prompts must not be flagged as token-id prompts.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + data = {"prompt": empty_prompt} + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), + cache=None, + data=data, + call_type="text_completion", + ) + + mock_check.assert_not_called() + assert result is data + assert result["prompt"] == empty_prompt + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "raw_prompt", + [ + [[1, 2, 3]], + [[1, 2], [3, 4]], + ["benign text", [99, 100]], + ], + ) + async def test_nested_token_id_prompt_raises_in_blocking_mode(self, raw_prompt): + """Nested/mixed token-id prompts must also be rejected in blocking pre_call mode.""" + guardrail = _make_guardrail() + + with patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), + cache=None, + data={"prompt": raw_prompt}, + call_type="text_completion", + ) + + mock_check.assert_not_called() + assert exc_info.value.status_code == 400 + assert "Token-id" in str(exc_info.value.detail) + + +class TestIsTokenIdPrompt: + @pytest.mark.parametrize( + "prompt,expected", + [ + ([1, 2, 3], True), + ([[1, 2, 3]], True), + ([[1, 2], [3, 4]], True), + (["hi", [1, 2]], True), + (["a", "b"], False), + ([], False), + ("hello", False), + (None, False), + ], + ) + def test_is_token_id_prompt(self, prompt, expected): + assert PurviewGuardrailBase.is_token_id_prompt(prompt) is expected + + +# --------------------------------------------------------------- +# Streaming iterator hook +# --------------------------------------------------------------- + + +class TestStreamingIteratorHook: + @pytest.mark.asyncio + async def test_streaming_clean_response_yields_all_chunks(self): + """Clean stream: all chunks must be re-yielded after DLP passes.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + + assembled_response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message(content="safe response", role="assistant"), + ) + ] + ) + + async def fake_response_stream(): + yield assembled_response + + with ( + patch("litellm.main.stream_chunk_builder", return_value=assembled_response), + patch( + "litellm.llms.base_llm.base_model_iterator.MockResponseIterator" + ) as mock_iterator_cls, + patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check, + ): + mock_check.return_value = {"policyActions": []} + + async def _iter_chunks(): + yield assembled_response + + mock_iterator_cls.return_value.__aiter__ = lambda s: _iter_chunks() + + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=fake_response_stream(), + request_data={"metadata": {"user_id": "user-123"}}, + ): + chunks.append(chunk) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["activity"] == "downloadText" + assert len(chunks) > 0 + + @pytest.mark.asyncio + async def test_streaming_violation_raises_before_any_chunk(self): + """A policy violation must raise HTTPException before yielding any chunk.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + + assembled_response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message( + content="SSN: 123-45-6789", + role="assistant", + ), + ) + ] + ) + + async def fake_response_stream(): + yield assembled_response + + with ( + patch("litellm.main.stream_chunk_builder", return_value=assembled_response), + patch.object( + guardrail, + "_check_content", + new_callable=AsyncMock, + side_effect=HTTPException( + status_code=400, + detail={ + "error": "Microsoft Purview DLP: Content blocked by policy" + }, + ), + ), + ): + chunks = [] + with pytest.raises(HTTPException) as exc_info: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key="test", user_id="user-123" + ), + response=fake_response_stream(), + request_data={"metadata": {"user_id": "user-123"}}, + ): + chunks.append(chunk) + + assert exc_info.value.status_code == 400 + assert len(chunks) == 0 # No chunks yielded before the block + + @pytest.mark.asyncio + async def test_streaming_no_user_id_raises_before_yield(self): + """No resolvable user_id → fail closed before any chunk is yielded.""" + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail = _make_guardrail() + + assembled_response = ModelResponse( + choices=[ + Choices( + index=0, + message=Message(content="some content", role="assistant"), + ) + ] + ) + + async def fake_response_stream(): + yield assembled_response + + with patch( + "litellm.main.stream_chunk_builder", return_value=assembled_response + ): + chunks = [] + with pytest.raises(HTTPException) as exc_info: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), # no user_id + response=fake_response_stream(), + request_data={}, + ): + chunks.append(chunk) + + assert exc_info.value.status_code == 400 + assert len(chunks) == 0 + + @pytest.mark.asyncio + async def test_streaming_text_completion_scans_before_yield(self): + """Streamed /v1/completions must be DLP-scanned via TextCompletionResponse.""" + from litellm.types.utils import TextChoices, TextCompletionResponse + + guardrail = _make_guardrail() + + assembled_response = TextCompletionResponse( + model="gpt-3.5-turbo-instruct", + choices=[TextChoices(text="completion body", index=0)], + ) + + async def fake_response_stream(): + yield assembled_response + + with ( + patch("litellm.main.stream_chunk_builder", return_value=assembled_response), + patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check, + ): + mock_check.return_value = {"policyActions": []} + + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=fake_response_stream(), + request_data={}, + ): + chunks.append(chunk) + + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["text"] == "completion body" + assert len(chunks) > 0 + + @pytest.mark.asyncio + async def test_streaming_responses_api_scans_completed_event(self): + """Streamed Responses API: assembled ResponsesAPIResponse must be DLP-scanned.""" + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + guardrail = _make_guardrail() + + completed_response = ResponsesAPIResponse( + id="resp-stream", + created_at=0, + output=[ + { + "type": "message", + "id": "msg-stream", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "streamed output"}], + } + ], + ) + created_event = ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=completed_response, + ) + completed_event = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=completed_response, + ) + + async def fake_response_stream(): + yield created_event + yield completed_event + + with ( + patch("litellm.main.stream_chunk_builder") as mock_stream_builder, + patch.object( + guardrail, "_check_content", new_callable=AsyncMock + ) as mock_check, + ): + mock_check.return_value = {"policyActions": []} + + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=fake_response_stream(), + request_data={}, + ): + chunks.append(chunk) + + mock_stream_builder.assert_not_called() + mock_check.assert_called_once() + assert mock_check.call_args.kwargs["activity"] == "downloadText" + assert mock_check.call_args.kwargs["text"] == "streamed output" + assert chunks == [created_event, completed_event] + + @pytest.mark.asyncio + async def test_streaming_responses_api_violation_blocks_before_yield(self): + """Responses API stream with a DLP violation must raise before any chunk is yielded.""" + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + guardrail = _make_guardrail() + + completed_response = ResponsesAPIResponse( + id="resp-stream-block", + created_at=0, + output=[ + { + "type": "message", + "id": "msg-stream-block", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "SSN: 123-45-6789"}], + } + ], + ) + completed_event = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=completed_response, + ) + + async def fake_response_stream(): + yield completed_event + + with patch.object( + guardrail, + "_check_content", + new_callable=AsyncMock, + side_effect=HTTPException( + status_code=400, + detail={"error": "Microsoft Purview DLP: Content blocked by policy"}, + ), + ): + chunks = [] + with pytest.raises(HTTPException) as exc_info: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key="test", user_id="user-123" + ), + response=fake_response_stream(), + request_data={}, + ): + chunks.append(chunk) + + assert exc_info.value.status_code == 400 + assert len(chunks) == 0 + + +# --------------------------------------------------------------- +# Auto-discovery registration +# --------------------------------------------------------------- + + +class TestRegistration: + def test_registry_contains_microsoft_purview(self): + from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview import ( + guardrail_class_registry, + guardrail_initializer_registry, + ) + + assert "microsoft_purview" in guardrail_initializer_registry + assert "microsoft_purview" in guardrail_class_registry + assert ( + guardrail_class_registry["microsoft_purview"] + is MicrosoftPurviewDLPGuardrail + )