diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/__init__.py new file mode 100644 index 00000000000..813cf1bc57d --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/__init__.py @@ -0,0 +1,73 @@ +"""Alice WonderFence guardrail integration for LiteLLM.""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .alice_wonderfence import ( + WonderFenceBlockedError, + WonderFenceGuardrail, + WonderFenceMissingSecrets, +) + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> WonderFenceGuardrail: + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Alice WonderFence guardrail requires a guardrail_name") + + # Pass only fields the user (or pydantic default) actually populated. The + # constructor owns the defaults, so `or X` chains here would silently + # override explicit falsy values like `api_timeout=0` or `fail_open=False`. + init_kwargs: dict = { + "guardrail_name": guardrail_name, + "api_key": litellm_params.api_key, + "api_base": litellm_params.api_base, + "platform": litellm_params.platform, + "max_cached_clients": litellm_params.max_cached_clients, + "connection_pool_limit": litellm_params.connection_pool_limit, + "event_hook": litellm_params.mode, + "default_on": (litellm_params.default_on if litellm_params.default_on is not None else True), + } + if litellm_params.api_timeout is not None: + init_kwargs["api_timeout"] = litellm_params.api_timeout + if litellm_params.fail_open is not None: + init_kwargs["fail_open"] = litellm_params.fail_open + if litellm_params.block_message is not None: + init_kwargs["block_message"] = litellm_params.block_message + if litellm_params.debug is not None: + init_kwargs["debug"] = litellm_params.debug + if litellm_params.allow_request_metadata_override is not None: + init_kwargs["allow_request_metadata_override"] = litellm_params.allow_request_metadata_override + if litellm_params.max_scan_chars is not None: + init_kwargs["max_scan_chars"] = litellm_params.max_scan_chars + if litellm_params.max_scan_segments is not None: + init_kwargs["max_scan_segments"] = litellm_params.max_scan_segments + + wonderfence_guardrail = WonderFenceGuardrail(**init_kwargs) + + litellm.logging_callback_manager.add_litellm_callback(wonderfence_guardrail) + return wonderfence_guardrail + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.ALICE_WONDERFENCE.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.ALICE_WONDERFENCE.value: WonderFenceGuardrail, +} + + +__all__ = [ + "WonderFenceBlockedError", + "WonderFenceGuardrail", + "WonderFenceMissingSecrets", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/alice_wonderfence.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/alice_wonderfence.py new file mode 100644 index 00000000000..8d6e328d449 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/alice_wonderfence.py @@ -0,0 +1,442 @@ +"""Alice WonderFence guardrail integration for LiteLLM.""" + +import logging +import os +from collections import OrderedDict +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any, Literal, Optional + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, +) +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.proxy.guardrails.guardrail_hooks.alice_wonderfence import ( + WonderFenceGuardrailConfigModel, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +from .chunked_evaluation import ( + DEFAULT_MAX_CONCURRENCY, + evaluate_segments, +) +from .client_cache import ClientBuildSpec, get_or_create_client, load_sdk +from .credentials import CredentialConfig, resolve_credentials +from .exceptions import ( + WonderFenceBlockedError, + WonderFenceMissingSecrets, + WonderFenceScanBudgetExceeded, +) +from .processing import ( + JOINER, + apply_response_verdicts, + block_detail, + build_analysis_context, + check_scan_budget, + function_definition_segments, + raise_if_blocked, + reconstruct, + tool_call_arg_segments, + tool_definition_segments, +) + +if TYPE_CHECKING: + from wonderfence_sdk.client import ( # pyright: ignore[reportMissingTypeStubs] # wonderfence_sdk ships no type stubs + WonderFenceV2Client as _WonderFenceV2Client, + ) + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.base import ( + GuardrailConfigModel, + ) + + +logger = verbose_proxy_logger.getChild("alice_wonderfence") + + +def _is_wonderfence_client_error(exc: Exception) -> bool: + """True for a persistent WonderFence client/config error. + + The V2 SDK raises ``Exception(":")`` on an HTTP error, so an + invalid or revoked api_key / app_id surfaces as a 4xx. Such failures are + persistent, not a transient outage, so they must never fail open: doing so + would silently disable scanning for every affected request until the + credential is fixed. 429 and 5xx stay in the fail-open path, matching the + SDK's own retry classification (``utils.py``: 4xx except 429 is not retried). + """ + head = str(exc).split(":", 1)[0].strip() + if not head.isdigit(): + return False + return 400 <= int(head) < 500 and int(head) != 429 + + +class WonderFenceGuardrail(CustomGuardrail): + """Alice WonderFence guardrail handler using the V2 SDK client. + + ``api_key`` and ``app_id`` are resolved per request from API-key metadata, + team metadata, optionally request metadata, with ``api_key`` falling back + to a configured default. ``app_id`` has no default. See ``credentials`` + module for the full precedence rationale. + + A V2 SDK client is cached per resolved ``api_key`` (LRU). + """ + + def __init__( + self, + guardrail_name: str, + api_key: str | None = None, + api_base: str | None = None, + api_timeout: float = 10.0, + platform: str | None = None, + fail_open: bool = False, + block_message: str = "Content violates our policies and has been blocked", + debug: bool = False, + max_cached_clients: int | None = None, + connection_pool_limit: int | None = None, + allow_request_metadata_override: bool = False, + max_scan_chars: int | None = None, + max_scan_segments: int | None = None, + event_hook: (GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None) = None, + default_on: bool = True, + **kwargs: Any, + ) -> None: + """Initialize the Alice WonderFence guardrail. + + Args: + guardrail_name: Unique identifier for this guardrail instance. + api_key: Default WonderFence API key. Falls back to ``ALICE_API_KEY``. + api_base: Optional base URL override for the WonderFence API. + api_timeout: Per-call timeout in seconds (rounded to int for SDK). + platform: Cloud platform identifier (e.g., aws, azure, databricks). + fail_open: When True, allow requests/responses through if WonderFence + is unreachable. BLOCK actions are always enforced. + block_message: User-facing error message returned on BLOCK action. + debug: Set guardrail logger to DEBUG level. + max_cached_clients: Max SDK clients cached per guardrail (LRU, + keyed by api_key). Default 10. Env: ALICE_MAX_CACHED_CLIENTS. + connection_pool_limit: Max connections per SDK client HTTP pool. + Env: ALICE_CONNECTION_POOL_LIMIT. + allow_request_metadata_override: When True, allow per-request + ``metadata.alice_wonderfence_api_key`` / + ``metadata.alice_wonderfence_app_id`` as a last-resort source + (after API-key and team metadata). Defaults to False so + caller-controlled fields cannot bypass admin-pinned credentials. + max_scan_chars: Fail-closed total-work cap on combined scan + characters per request/response. Default 1_000_000. Env: + ALICE_MAX_SCAN_CHARS. + max_scan_segments: Fail-closed total-work cap on scan segment count + per request/response. Default 1_000. Env: + ALICE_MAX_SCAN_SEGMENTS. + event_hook: Event hook mode. + default_on: Whether the guardrail is enabled by default. + """ + WonderFenceV2Client, AnalysisContext = load_sdk() + self._WonderFenceV2Client = WonderFenceV2Client + self._AnalysisContext = AnalysisContext + + self.api_key = api_key or os.environ.get("ALICE_API_KEY") + self.api_base = api_base + self.api_timeout = api_timeout + self.platform = platform + self.fail_open = fail_open + self.block_message = block_message + self.allow_request_metadata_override = allow_request_metadata_override + env_max_chars = os.environ.get("ALICE_MAX_SCAN_CHARS") + self.max_scan_chars: int | None = ( + max_scan_chars if max_scan_chars is not None else (int(env_max_chars) if env_max_chars else 1_000_000) + ) + env_max_segments = os.environ.get("ALICE_MAX_SCAN_SEGMENTS") + self.max_scan_segments: int | None = ( + max_scan_segments + if max_scan_segments is not None + else (int(env_max_segments) if env_max_segments else 1_000) + ) + + if debug: + logger.setLevel(logging.DEBUG) + + self._client_cache: OrderedDict[str, _WonderFenceV2Client] = OrderedDict() + self._client_cache_maxsize = max_cached_clients or int(os.environ.get("ALICE_MAX_CACHED_CLIENTS", "10")) + env_pool = os.environ.get("ALICE_CONNECTION_POOL_LIMIT") + self._connection_pool_limit: int | None = ( + connection_pool_limit if connection_pool_limit is not None else (int(env_pool) if env_pool else None) + ) + + supported_event_hooks = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ] + + super().__init__( + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + supported_event_hooks=supported_event_hooks, + **kwargs, + ) + # Narrow attribute type: base class declares Optional[str], but our + # __init__ requires a non-empty string and the factory rejects empty. + self.guardrail_name: str = guardrail_name + + key_suffix = f"***{self.api_key[-4:]}" if self.api_key else "" + logger.debug( + "Alice WonderFence guardrail initialized: name=%s default_api_key=%s", + guardrail_name, + key_suffix, + ) + + async def _get_client(self, api_key: str) -> "_WonderFenceV2Client": + """Return a cached WonderFenceV2Client for the given api_key (LRU).""" + return get_or_create_client( + api_key, + self._client_cache, + self._client_cache_maxsize, + ClientBuildSpec( + client_class=self._WonderFenceV2Client, + api_timeout=self.api_timeout, + api_base=self.api_base, + connection_pool_limit=self._connection_pool_limit, + ), + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Apply WonderFence guardrail using V2 client + per-request app_id. + + Request side joins all scan pieces (message text plus detection-only + tool-call args and tool/function descriptions) into one document and + scans it in ~1 call (chunked only when it exceeds the size limit), so + call volume scales with total size rather than message count. Response + side stays per-segment (independent choices / model tool-call args) + because the handler's response write-back is purely positional and has + no ``structured_messages`` path. + """ + texts = inputs.get("texts") or [] + tool_indices, tool_arg_segments = tool_call_arg_segments(inputs) + tool_def_texts = tool_definition_segments(inputs) + # Legacy top-level functions[] only exist on the request body; the + # translation layer does not surface them in inputs, so read request_data. + function_def_texts = function_definition_segments(request_data) if input_type == "request" else [] + + if input_type == "request": + scan_pieces = [*texts, *tool_arg_segments, *tool_def_texts, *function_def_texts] + else: + scan_pieces = [*texts, *tool_arg_segments] + + if not scan_pieces: + logger.debug( + "Alice WonderFence (apply_guardrail): nothing to scan for %s", + input_type, + ) + return inputs + + try: + check_scan_budget(scan_pieces, self.max_scan_chars, self.max_scan_segments) + api_key, app_id = resolve_credentials( + request_data, + input_type, + logging_obj, + CredentialConfig( + guardrail_name=self.guardrail_name, + default_api_key=self.api_key, + allow_request_metadata_override=self.allow_request_metadata_override, + ), + ) + client = await self._get_client(api_key) + context = build_analysis_context(request_data, self.platform, self._AnalysisContext) + max_concurrency = self._connection_pool_limit or DEFAULT_MAX_CONCURRENCY + + if input_type == "request": + + async def evaluate(text: str) -> object: + return await client.evaluate_prompt(app_id=app_id, prompt=text, context=context, custom_fields=None) + + await self._scan_request(inputs, scan_pieces, len(texts), evaluate, max_concurrency, app_id) + else: + + async def evaluate(text: str) -> object: + return await client.evaluate_response( + app_id=app_id, + response=text, + context=context, + custom_fields=None, + ) + + await self._scan_response(inputs, texts, tool_indices, tool_arg_segments, evaluate, max_concurrency) + + except WonderFenceScanBudgetExceeded as e: + # Fail-closed config/abuse guard: reject before any provider call and + # never fall through to the fail_open path below. + raise HTTPException(status_code=400, detail=e.detail) + except WonderFenceBlockedError as e: + raise HTTPException(status_code=400, detail=e.detail) + except WonderFenceMissingSecrets as e: + # Configuration errors (no api_key / app_id resolvable) are never + # fail-open: a misconfigured tenant must not silently bypass the + # guardrail. + raise HTTPException( + status_code=500, + detail={ + "error": "Error in Alice WonderFence Guardrail", + "guardrail_name": self.guardrail_name, + "exception": str(e), + }, + ) from e + except Exception as e: + if self.fail_open and not _is_wonderfence_client_error(e): + # Log only — do not add to the applied-guardrails header. The + # header lists configured guardrail_names verbatim; consumers + # rely on its membership to decide whether scanning ran. A + # synthetic suffix (e.g. ":unscanned") would silently pass the + # membership check and mask audit gaps. + logger.error( + "Alice WonderFence unreachable; fail-open enabled, proceeding " + "without guardrail. guardrail_name=%s input_type=%s " + "guardrail_status=unscanned_fail_open error=%s", + self.guardrail_name, + input_type, + str(e), + exc_info=e, + ) + return inputs + # Either fail-open is off, or this is a persistent client/config + # error (4xx) that must never fail open — see + # _is_wonderfence_client_error. + logger.error( + "Alice WonderFence request failed and was not allowed through " + "(fail-open off or persistent client error). guardrail_name=%s input_type=%s error=%s", + self.guardrail_name, + input_type, + str(e), + exc_info=e, + ) + raise HTTPException( + status_code=500, + detail={ + "error": "Error in Alice WonderFence Guardrail", + "guardrail_name": self.guardrail_name, + "exception": str(e), + }, + ) from e + + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) + return inputs + + async def _scan_request( + self, + inputs: GenericGuardrailAPIInputs, + pieces: list[str], + n_text: int, + evaluate: "Callable[[str], Awaitable[object]]", + max_concurrency: int, + app_id: str, + ) -> None: + """Scan the joined request document and write MASK back to ``texts``. + + The first ``n_text`` pieces are maskable message-text parts; the rest are + detection-only tool-call args and tool/function descriptions. On MASK we + reconstruct per-part masked text by aligning the join against the masked + document and write the recovered message-text parts back to + ``inputs["texts"]`` (positional write-back / "Path B"): the handler + already maps that list onto the right message parts, so there is no need + to rebuild ``structured_messages``. Reconstruction failure fails closed + (block) rather than misassigning a redaction. + """ + document = JOINER.join(pieces) + logger.debug( + "Alice WonderFence (apply_guardrail request): scanning joined document of %d piece(s) " + "(%d text + %d detection-only), %d chars, guardrail=%s app_id=%s", + len(pieces), + n_text, + len(pieces) - n_text, + len(document), + self.guardrail_name, + app_id, + ) + verdict = (await evaluate_segments([document], evaluate, max_concurrency=max_concurrency))[0] + raise_if_blocked([verdict], self.guardrail_name, self.block_message) + + correlation_id = verdict.correlation_ids[0] if verdict.correlation_ids else None + if verdict.action == "MASK": + recovered = reconstruct(pieces, verdict.masked_chunks) + if recovered is None: + logger.warning( + "Alice WonderFence (apply_guardrail request): MASK reconstruction unavailable " + "(document too large, or a joiner / part boundary landed inside a masked span); " + "failing closed. guardrail=%s correlation_id=%s", + self.guardrail_name, + correlation_id, + ) + raise WonderFenceBlockedError(block_detail([verdict], self.guardrail_name, self.block_message)) + if recovered[n_text:] != pieces[n_text:]: + # The mask redacted a detection-only piece (tool-call args or a + # tool / function description). Those are not maskable in place + # (the joined form is not the wire format), so we cannot forward + # the original unredacted value; fail closed rather than leak it. + logger.warning( + "Alice WonderFence (apply_guardrail request): MASK landed in a detection-only " + "piece (tool-call args / tool or function description); failing closed. guardrail=%s correlation_id=%s", + self.guardrail_name, + correlation_id, + ) + raise WonderFenceBlockedError(block_detail([verdict], self.guardrail_name, self.block_message)) + inputs["texts"] = recovered[:n_text] + logger.info( + "Alice WonderFence (apply_guardrail request): MASK applied to request text guardrail=%s correlation_id=%s", + self.guardrail_name, + correlation_id, + ) + elif verdict.action == "DETECT": + logger.warning( + "Alice WonderFence (apply_guardrail request): DETECT on joined document guardrail=%s correlation_id=%s", + self.guardrail_name, + correlation_id, + ) + + async def _scan_response( + self, + inputs: GenericGuardrailAPIInputs, + texts: list[str], + tool_indices: list[int], + tool_arg_segments: list[str], + evaluate: "Callable[[str], Awaitable[object]]", + max_concurrency: int, + ) -> None: + """Scan response segments per-index and write masks back in place.""" + segments = [*texts, *tool_arg_segments] + logger.debug( + "Alice WonderFence (apply_guardrail response): evaluating %d text + %d tool-call segment(s) guardrail=%s", + len(texts), + len(tool_arg_segments), + self.guardrail_name, + ) + verdicts = await evaluate_segments(segments, evaluate, max_concurrency=max_concurrency) + n_text = len(texts) + apply_response_verdicts( + inputs, + verdicts[:n_text], + tool_indices, + verdicts[n_text:], + self.guardrail_name, + self.block_message, + ) + + @staticmethod + def get_config_model() -> type["GuardrailConfigModel"] | None: + """Return the config model for UI rendering.""" + return WonderFenceGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/chunked_evaluation.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/chunked_evaluation.py new file mode 100644 index 00000000000..444aaff89d8 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/chunked_evaluation.py @@ -0,0 +1,178 @@ +"""Chunk, evaluate in parallel, and aggregate per segment. + +Guardrail-agnostic: the only coupling to WonderFence is the injected +``evaluate`` callable and the result shape it returns (``action``, +``action_text``, ``detections``, ``correlation_id``). Replace ``evaluate`` to +target a different backend. +""" + +import asyncio +import re +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any + +MAX_PROMPT_CHARS = 10000 # WonderFence server-side prompt limit +DEFAULT_MAX_CONCURRENCY = 10 # used when the client connection_pool_limit is unset +# Overlap: a segment longer than the prompt limit is split into disjoint "owned" +# regions, but each chunk is scanned with the last N chars of the previous owned +# region prepended as a read-only prefix. A phrase straddling an owned-region +# seam (up to N chars into the left region) is therefore seen whole by one scan, +# so it can BLOCK/DETECT and, when the service masks it, the prefix bytes change +# and we fail closed (see ``_aggregate``) rather than stitch a half-masked seam. +# This replaces the old separate boundary-window calls. Confirm sizing with the +# WonderFence team alongside MAX_PROMPT_CHARS. +CHUNK_OVERLAP_CHARS = 512 + + +@dataclass +class SegmentVerdict: + action: str # "BLOCK" | "MASK" | "DETECT" | "" + masked_text: str | None + detections: list + correlation_ids: list[str] + # Per owned-region ``(original, masked)`` pairs, set only on MASK. Lets the + # caller align masking back to sub-structure (e.g. joined message parts) one + # chunk at a time instead of over the whole document, so the alignment cost + # is bounded by the chunk size rather than quadratic in the segment length. + masked_chunks: list[tuple[str, str]] | None = None + + +def _split_text(text: str, max_chars: int) -> list[str]: + """Split ``text`` into <= ``max_chars`` chunks with ``"".join(chunks) == text``. + + Splits at whitespace boundaries; whitespace runs are preserved as their own + tokens so the rejoin is byte-identical. A single token longer than + ``max_chars`` is force-split. Always returns at least one chunk. + """ + if len(text) <= max_chars: + return [text] + + tokens = re.findall(r"\S+|\s+", text) + chunks: list[str] = [] + current = "" + for token in tokens: + if len(current) + len(token) <= max_chars: + current += token + continue + if current: + chunks.append(current) + current = "" + while len(token) > max_chars: + chunks.append(token[:max_chars]) + token = token[max_chars:] + current = token + if current: + chunks.append(current) + return chunks + + +def _action_str(result: object) -> str: + action = getattr(result, "action", "") + return action.value if hasattr(action, "value") else (action or "") + + +def _overlap_chunks(text: str, max_chars: int, overlap: int) -> list[tuple[str, str]]: + """Split ``text`` into overlapping scan chunks as ``(prefix, owned)`` pairs. + + ``owned`` regions are disjoint and concatenate back to ``text`` (lossless); + ``prefix`` is the last ``overlap`` chars of the previous owned region (empty + for the first). The scan input for a chunk is ``prefix + owned``, giving + ``overlap`` chars of left-context so a phrase straddling the owned-region + seam is seen whole. Reassembly strips the verbatim prefix back off, so the + owned regions still rejoin losslessly. + """ + owned = _split_text(text, max(1, max_chars - overlap)) + return [(owned[i - 1][-overlap:] if i and overlap > 0 else "", region) for i, region in enumerate(owned)] + + +def _aggregate(chunks: list[tuple[str, str]], results: list[Any]) -> SegmentVerdict: + """Fold per-chunk results (scans of ``prefix + owned``) into one verdict. + + Precedence BLOCK > MASK > DETECT > NO_ACTION. A MASK whose masked text no + longer starts with its verbatim ``prefix`` means the redaction reached into + the prefix bytes -- i.e. content straddling (or sitting within ``overlap`` of) + the owned-region seam. That cannot be stitched back without double-counting + the overlap, so it fails closed (BLOCK) rather than leak the un-redacted half. + Otherwise the prefix is stripped by its known length (no alignment needed) + and the owned regions rejoin into the masked segment; the per-chunk + ``(original, masked)`` pairs are carried on the verdict for bounded caller-side + alignment. + """ + actions = [_action_str(r) for r in results] + detections: list = [] + correlation_ids: list[str] = [] + for r in results: + detections.extend(getattr(r, "detections", None) or []) + cid = getattr(r, "correlation_id", None) + if cid: + correlation_ids.append(cid) + + if "BLOCK" in actions: + return SegmentVerdict("BLOCK", None, detections, correlation_ids) + + if "MASK" in actions: + masked_chunks: list[tuple[str, str]] = [] + for (prefix, owned), r in zip(chunks, results): + if _action_str(r) != "MASK": + masked_chunks.append((owned, owned)) + continue + masked = r.action_text if getattr(r, "action_text", None) is not None else prefix + "[MASKED]" + if not masked.startswith(prefix): + return SegmentVerdict("BLOCK", None, detections, correlation_ids) + masked_chunks.append((owned, masked[len(prefix) :])) + masked_text = "".join(m for _, m in masked_chunks) + return SegmentVerdict("MASK", masked_text, detections, correlation_ids, masked_chunks) + + if "DETECT" in actions: + return SegmentVerdict("DETECT", None, detections, correlation_ids) + return SegmentVerdict("", None, detections, correlation_ids) + + +async def evaluate_segments( + segments: list[str], + evaluate: Callable[[str], Awaitable[Any]], + max_chars: int = MAX_PROMPT_CHARS, + max_concurrency: int = DEFAULT_MAX_CONCURRENCY, + overlap: int = CHUNK_OVERLAP_CHARS, +) -> list[SegmentVerdict]: + """Evaluate every segment (chunked) in parallel; return one verdict per segment. + + Each segment is split into <= ``max_chars`` overlapping chunks (disjoint + ``owned`` regions each carrying an ``overlap``-char read-only prefix from the + previous region, see ``_overlap_chunks``) so a phrase straddling an + owned-region seam is seen whole by one scan without a separate boundary call. + Every chunk across every segment is evaluated through a single + ``asyncio.gather`` behind one shared ``Semaphore(max_concurrency)``. Results + are folded per segment (see ``_aggregate``) with precedence + BLOCK > MASK > DETECT > NO_ACTION. + + The request side passes a single joined document here (one segment) so the + common case is one call; the response side passes one segment per choice / + tool-call arg. There is no cross-segment overlap: message parts are already + joined into one document on the request side, and response choices are + independent. + """ + semaphore = asyncio.Semaphore(max_concurrency) + + async def run(text: str) -> object: + async with semaphore: + return await evaluate(text) + + # Keep each chunk's scan input (prefix + owned) within the prompt limit. + ov = min(overlap, max_chars // 2) + seg_chunks = [_overlap_chunks(s, max_chars, ov) for s in segments] + + index: list[tuple[int, int]] = [] + tasks = [] + for si, chunks in enumerate(seg_chunks): + for ci, (prefix, owned) in enumerate(chunks): + index.append((si, ci)) + tasks.append(run(prefix + owned)) + results = await asyncio.gather(*tasks) + + chunk_res: list[list[Any]] = [[None] * len(c) for c in seg_chunks] + for (si, ci), res in zip(index, results): + chunk_res[si][ci] = res + + return [_aggregate(seg_chunks[si], chunk_res[si]) for si in range(len(segments))] diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/client_cache.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/client_cache.py new file mode 100644 index 00000000000..6b6ee13dd99 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/client_cache.py @@ -0,0 +1,80 @@ +"""WonderFence SDK loader + per-api_key LRU client cache.""" + +from collections import OrderedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from wonderfence_sdk.client import ( # pyright: ignore[reportMissingTypeStubs] # wonderfence_sdk ships no type stubs + WonderFenceV2Client as _WonderFenceV2Client, + ) + + +@dataclass(frozen=True) +class ClientBuildSpec: + """How to construct a WonderFenceV2Client on a cache miss. + + ``platform`` is deliberately absent: it is a per-request analysis attribute + that belongs on ``AnalysisContext`` (set by ``build_analysis_context``), not + on the client constructor. The V2 client signature is + ``(api_key, base_url, *, api_timeout, connection_pool_limit)`` with no + ``platform`` parameter, so forwarding it here raised ``TypeError`` on every + scan. + """ + + client_class: Callable[..., object] + api_timeout: float + api_base: str | None + connection_pool_limit: int | None + + +def load_sdk() -> tuple[Any, Any]: + """Lazy-import WonderFence SDK classes (``WonderFenceV2Client``, ``AnalysisContext``). + + Deferred to instance construction (not module load) because wonderfence_sdk + is an optional dependency: importing it at module top would break litellm + installs that don't use this guardrail. Callers cache the returned classes + on the instance so per-call hot paths don't re-trigger the import machinery. + """ + try: + from wonderfence_sdk.client import ( # pyright: ignore[reportMissingTypeStubs] # wonderfence_sdk ships no type stubs + WonderFenceV2Client, + ) + from wonderfence_sdk.models import ( # pyright: ignore[reportMissingTypeStubs] # wonderfence_sdk ships no type stubs + AnalysisContext, + ) + except ImportError as e: + raise ImportError("Alice WonderFence SDK not installed. Install with: pip install wonderfence-sdk") from e + return WonderFenceV2Client, AnalysisContext + + +def get_or_create_client( + api_key: str, + cache: "OrderedDict[str, _WonderFenceV2Client]", + cache_maxsize: int, + spec: ClientBuildSpec, +) -> "_WonderFenceV2Client": + """LRU client lookup keyed by ``api_key``; construct on miss.""" + if api_key in cache: + cache.move_to_end(api_key) + return cache[api_key] + + client_kwargs: dict = { + "api_key": api_key, + "api_timeout": round(spec.api_timeout), + } + if spec.api_base: + client_kwargs["base_url"] = spec.api_base + if spec.connection_pool_limit is not None: + client_kwargs["connection_pool_limit"] = spec.connection_pool_limit + + client = spec.client_class(**client_kwargs) + cache[api_key] = client + + if len(cache) > cache_maxsize: + # Drop reference only — never close. An evicted client may still be + # held by in-flight apply_guardrail coroutines; closing it would + # break their pooled HTTP connections. GC handles cleanup. + cache.popitem(last=False) + + return client diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/credentials.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/credentials.py new file mode 100644 index 00000000000..c540b1b0e5b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/credentials.py @@ -0,0 +1,266 @@ +"""Credential resolution + request-scoped stash for Alice WonderFence. + +Resolves ``api_key`` / ``app_id`` per request from API-key metadata, team +metadata, optionally request metadata, with ``api_key`` falling back to a +configured default. ``app_id`` has no default. + +Admin-pinned credentials (key/team metadata) always win over request metadata +so a caller cannot bypass their assigned WonderFence app. +``allow_request_metadata_override`` defaults to False; enable only for +trusted-gateway deployments that need request-level overrides. + +The two metadata buckets (``metadata`` and ``litellm_metadata``) are merged +with the proxy-injected ``litellm_metadata`` winning on key collision, so admin +pins cannot be shadowed by a caller-supplied ``metadata`` body — see +``get_metadata``. + +The stash bridges pre_call resolution into post_call where request metadata is +gone — see ``stash_resolved`` for the full rationale. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, Optional + +from .exceptions import WonderFenceMissingSecrets + + +@dataclass(frozen=True) +class CredentialConfig: + """Per-guardrail config consulted during credential resolution.""" + + guardrail_name: str + default_api_key: str | None + allow_request_metadata_override: bool + + +def _nonempty_str(value: object) -> str | None: + """Return ``value`` only if it is a non-empty/non-blank string, else None. + + Credential sources (request body, key/team metadata, config default) are + only honored when they carry a real string. A truthy non-string override + (list, dict, number) must not pass through to the SDK, where it would raise + a type error that ``fail_open`` could swallow into a skipped scan. + """ + return value if isinstance(value, str) and value.strip() else None + + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + + +# Prefix for the per-guardrail attribute on the request-scoped logging_obj where +# the resolved (api_key, app_id) is stashed so post_call can recover it. The +# guardrail name is baked into the attribute so each instance has a physically +# separate slot — one instance can never read another's credentials. +# +# These are private instance attributes, NOT model_call_details keys, because +# model_call_details is forwarded verbatim as ``kwargs`` to every logging +# callback / exporter (litellm_logging.py) — stashing the resolved WonderFence +# api_key there would leak a tenant secret into logs. A private attribute on the +# same object gives the identical cross-hook / cross-task visibility (see +# stash_resolved) without entering the logged payload. +_STASH_ATTR_PREFIX = "_alice_wonderfence_resolved__" + + +def _stash_attr(guardrail_name: str) -> str: + return _STASH_ATTR_PREFIX + guardrail_name + + +def get_metadata(request_data: dict) -> dict: + """Merge caller metadata with proxy-injected litellm_metadata. + + Proxy-injected values win on key collision so admin-pinned + user_api_key_metadata / user_api_key_team_metadata can never be shadowed + by a caller-supplied `metadata` body. On routes in LITELLM_METADATA_ROUTES + (e.g. /v1/responses) the admin pins live in `litellm_metadata` while the + caller bucket is `metadata`; on /chat/completions they coincide. + """ + caller = request_data.get("metadata") + litellm_md = request_data.get("litellm_metadata") + caller = caller if isinstance(caller, dict) else {} + litellm_md = litellm_md if isinstance(litellm_md, dict) else {} + return {**caller, **litellm_md} + + +def resolve_api_key( + request_data: dict, + default_api_key: str | None, + allow_request_metadata_override: bool, +) -> str: + """Resolve api_key from key → team → (request, when opt-in) → default. + + Admin-pinned sources (API-key and team metadata) take precedence over + request-body metadata so a caller cannot bypass their assigned WonderFence + credentials. Request metadata is consulted only when + ``allow_request_metadata_override`` is True, and even then only after the + admin-controlled sources. + + The LiteLLM framework copies key/team metadata from ``UserAPIKeyAuth`` into + ``data['metadata']`` under ``user_api_key_metadata`` and + ``user_api_key_team_metadata``, so all sources are read from + ``request_data``. + """ + metadata = get_metadata(request_data) + + key_metadata = metadata.get("user_api_key_metadata") or {} + if isinstance(key_metadata, dict): + val = _nonempty_str(key_metadata.get("alice_wonderfence_api_key")) + if val: + return val + + team_metadata = metadata.get("user_api_key_team_metadata") or {} + if isinstance(team_metadata, dict): + val = _nonempty_str(team_metadata.get("alice_wonderfence_api_key")) + if val: + return val + + if allow_request_metadata_override: + val = _nonempty_str(metadata.get("alice_wonderfence_api_key")) + if val: + return val + + val = _nonempty_str(default_api_key) + if val: + return val + + raise WonderFenceMissingSecrets( + "No alice_wonderfence_api_key found in API-key metadata, team " + "metadata, request metadata (when allow_request_metadata_override " + "is enabled), or default config (ALICE_API_KEY)." + ) + + +def resolve_app_id(request_data: dict, allow_request_metadata_override: bool) -> str: + """Resolve app_id from key → team → (request, when opt-in). No default. + + Admin-pinned sources win over request-body metadata; request metadata is + only consulted when ``allow_request_metadata_override`` is True. Raises + ``WonderFenceMissingSecrets`` when nothing resolves. + """ + metadata = get_metadata(request_data) + + key_metadata = metadata.get("user_api_key_metadata") or {} + if isinstance(key_metadata, dict): + val = _nonempty_str(key_metadata.get("alice_wonderfence_app_id")) + if val: + return val + + team_metadata = metadata.get("user_api_key_team_metadata") or {} + if isinstance(team_metadata, dict): + val = _nonempty_str(team_metadata.get("alice_wonderfence_app_id")) + if val: + return val + + if allow_request_metadata_override: + val = _nonempty_str(metadata.get("alice_wonderfence_app_id")) + if val: + return val + + raise WonderFenceMissingSecrets( + "No alice_wonderfence_app_id found in API-key metadata, team " + "metadata, or request metadata (when allow_request_metadata_override " + "is enabled). app_id must be provided per request." + ) + + +def stash_resolved( + logging_obj: Optional["LiteLLMLoggingObj"], + guardrail_name: str, + api_key: str, + app_id: str, +) -> None: + """Persist resolved (api_key, app_id) on the request-scoped logging_obj + so post_call can recover it. + + Why we need this: + LiteLLM's per-provider chat translation handler synthesizes a fresh + ``request_data`` for post_call (``process_output_response``, e.g. + ``litellm/llms/openai/chat/guardrail_translation/handler.py:312``). + That dict only carries ``litellm_metadata.user_api_key_metadata`` and + ``user_api_key_team_metadata`` — the original request body's + ``metadata`` field (where per-request ``alice_wonderfence_app_id`` + lives) is dropped. Without a bridge, post_call resolution fails even + though the request explicitly supplied the value. + + Why a private attribute on logging_obj (and not model_call_details): + ``model_call_details`` is forwarded verbatim as ``kwargs`` to every + logging callback / exporter (``litellm_logging.py`` passes + ``kwargs=self.model_call_details`` to success/failure handlers and + logging hooks), and the redaction layer only scrubs message + input/output and known StandardLoggingPayload fields, not arbitrary + custom keys. Stashing the resolved WonderFence ``api_key`` there leaks + a tenant secret into logs. A private instance attribute is request + scoped on the same object but is not part of the logged ``kwargs``. + + Why an attribute on logging_obj (and not a ContextVar): + during_call hooks run via ``asyncio.gather`` in + ``litellm/proxy/utils.py:1500``, which wraps each coroutine in its own + asyncio Task with a *copied* context. ContextVar writes in a child + Task are not visible to the parent Task that runs post_call, so a + ContextVar bridge silently fails. ``logging_obj`` is passed through + every hook by reference (same object across pre_call, during_call, + and post_call), so mutations to it are visible regardless of task + boundary. + + The attribute is per-guardrail (see ``_stash_attr``) so multiple + alice_wonderfence instances on the same request each get an isolated slot. + """ + if logging_obj is None: + return + setattr(logging_obj, _stash_attr(guardrail_name), (api_key, app_id)) + + +def recover_resolved(logging_obj: Optional["LiteLLMLoggingObj"], guardrail_name: str) -> tuple[str, str] | None: + """Look up the (api_key, app_id) this guardrail stashed earlier in this + request, or ``None``. + + Returns only this instance's own stash. It deliberately does NOT fall back + to another alice_wonderfence instance's stash: a sibling may have resolved + under a different policy (e.g. ``allow_request_metadata_override=True``, + carrying caller-supplied request-body credentials) that a stricter instance + must not inherit. When this instance has no own stash, the caller fails + closed rather than borrowing. + """ + if logging_obj is None: + return None + value = getattr(logging_obj, _stash_attr(guardrail_name), None) + return value if isinstance(value, tuple) else None + + +def resolve_credentials( + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"], + config: CredentialConfig, +) -> tuple[str, str]: + """Resolve (api_key, app_id) for this call. + + For ``request``: read from request_data (canonical pre_call path) and stash + on logging_obj so post_call can recover. + + For ``response`` (post_call): try synthesized request_data first (works + when supplied via virtual key or team metadata, which the framework + preserves as ``litellm_metadata.user_api_key_metadata`` / + ``user_api_key_team_metadata``); fall back to the per-request logging_obj + stash for values supplied in the original request body's metadata, which + the framework drops before post_call. + """ + default_api_key = config.default_api_key + allow_override = config.allow_request_metadata_override + if input_type == "request": + api_key = resolve_api_key(request_data, default_api_key, allow_override) + app_id = resolve_app_id(request_data, allow_override) + stash_resolved(logging_obj, config.guardrail_name, api_key, app_id) + return api_key, app_id + try: + return ( + resolve_api_key(request_data, default_api_key, allow_override), + resolve_app_id(request_data, allow_override), + ) + except WonderFenceMissingSecrets: + recovered = recover_resolved(logging_obj, config.guardrail_name) + if recovered is None: + raise + return recovered diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/example_config.yaml b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/example_config.yaml new file mode 100644 index 00000000000..3254dc65a51 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/example_config.yaml @@ -0,0 +1,114 @@ +# Example LiteLLM Proxy configuration with Alice WonderFence guardrail +# +# Start the proxy with: +# litellm --config example_config.yaml +# +# Environment variables: +# ALICE_API_KEY - Default WonderFence API key (overridable per request) +# ALICE_MAX_CACHED_CLIENTS - Optional: max cached V2 SDK clients (default 10) +# ALICE_CONNECTION_POOL_LIMIT - Optional: HTTP pool size per client +# ALICE_MAX_SCAN_CHARS - Optional: total-work cap, max scan chars (default 1000000) +# ALICE_MAX_SCAN_SEGMENTS - Optional: total-work cap, max scan segments (default 1000) +# OPENAI_API_KEY - API key for OpenAI +# +# Per-key / per-team metadata keys (admin-controlled): +# alice_wonderfence_api_key - overrides default API key (optional) +# alice_wonderfence_app_id - REQUIRED — must be set on key or team +# +# Per-request metadata keys (caller-controlled, OFF by default): +# metadata.alice_wonderfence_api_key / metadata.alice_wonderfence_app_id are +# ignored unless allow_request_metadata_override is True on the guardrail. +# Even when enabled, key/team metadata still wins — request metadata is a +# last-resort source only. + +model_list: + - model_name: gpt-4 + litellm_params: + model: gpt-4 + api_key: os.environ/OPENAI_API_KEY + +guardrails: + + # Combined pre + post with advanced knobs + - guardrail_name: "alice-wonderfence" + litellm_params: + guardrail: alice_wonderfence + mode: ["pre_call", "post_call"] + api_key: os.environ/ALICE_API_KEY + api_timeout: 10.0 + platform: "aws" + default_on: false + debug: false + fail_open: false + max_cached_clients: 10 + block_message: "Content violates our policies and has been blocked by Alice WonderFence" + + # Every remaining message segment is evaluated (user, assistant, tool), + # not just the last user turn, so disallowed content placed in an earlier + # turn or an assistant prefill cannot slip past. System prompts are + # admin-controlled and excluded by default to avoid false positives; + # set this to false to scan them too. skip_tool_message_in_guardrail is + # the matching knob for tool messages. + skip_system_message_in_guardrail: true + + # Fail-closed total-work cap (DoS backstop). WonderFence has no batch API, + # so the request side joins all scan pieces into one document and scans it + # in ~1 call (chunked only past the size limit); these bounds reject an + # abusive request before any provider call and are never fail-open. + # max_scan_chars: 1000000 + # max_scan_segments: 1000 + + # connection_pool_limit: 20 + + # Enable only for trusted-gateway deployments that need to forward a + # per-tenant app_id/api_key from the request body. Even with this on, + # admin-pinned key/team metadata still wins; request metadata is a + # last-resort source only. Default is False — leave it off unless your + # callers are themselves trusted infrastructure. + allow_request_metadata_override: true + +# Example usage +# +# 1. Per-API-key app_id (set at key creation — recommended default): +# +# curl -X POST http://localhost:4000/key/generate \ +# -H "Authorization: Bearer sk-admin" \ +# -H "Content-Type: application/json" \ +# -d '{ +# "metadata": { +# "alice_wonderfence_app_id": "tenant-A-app", +# "alice_wonderfence_api_key": "wf-key-for-tenant-A" +# } +# }' +# +# 2. Per-team app_id (set at team creation): +# +# curl -X POST http://localhost:4000/team/new \ +# -H "Authorization: Bearer sk-admin" \ +# -H "Content-Type: application/json" \ +# -d '{ +# "team_alias": "team-billing", +# "metadata": { +# "alice_wonderfence_app_id": "team-billing-app" +# } +# }' +# +# 3. Request-level app_id (only when allow_request_metadata_override: true on +# the guardrail config — and even then, key/team metadata still wins): +# +# curl -X POST http://localhost:4000/chat/completions \ +# -H "Authorization: Bearer sk-xxx" \ +# -H "Content-Type: application/json" \ +# -d '{ +# "model": "gpt-4", +# "messages": [{"role": "user", "content": "Hello"}], +# "metadata": { +# "alice_wonderfence_app_id": "my-app-123", +# "session_id": "session-1" +# } +# }' +# +# Resolution priority (highest first): key metadata > team metadata > +# request metadata (only if allow_request_metadata_override=true) > +# config default. api_key falls back to config / ALICE_API_KEY env; +# app_id has NO default. diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/exceptions.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/exceptions.py new file mode 100644 index 00000000000..2ec9b85153c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/exceptions.py @@ -0,0 +1,27 @@ +"""Alice WonderFence guardrail exception types.""" + + +class WonderFenceMissingSecrets(Exception): + """Raised when Alice API key cannot be resolved from any source.""" + + +class WonderFenceBlockedError(Exception): + """Raised when WonderFence blocks a request/response.""" + + def __init__(self, detail: dict): + self.detail = detail + super().__init__(detail.get("error", "Blocked by Alice WonderFence guardrail")) + + +class WonderFenceScanBudgetExceeded(Exception): + """Raised when a request/response exceeds the configured total-work cap. + + A fail-closed configuration/abuse guard, never a transport failure: it is + mapped to HTTP 400 before any WonderFence call and is not subject to + ``fail_open`` (a caller must not be able to bypass scanning by overflowing + the cap). + """ + + def __init__(self, detail: dict): + self.detail = detail + super().__init__(detail.get("error", "Alice WonderFence scan budget exceeded")) diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/processing.py b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/processing.py new file mode 100644 index 00000000000..bf9e974afc0 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/processing.py @@ -0,0 +1,383 @@ +"""Pure transforms for Alice WonderFence: context build, scan-piece gathering, +joined-document masked reconstruction, response-side verdict apply, total-work cap.""" + +from collections.abc import Sequence +from difflib import SequenceMatcher +from itertools import accumulate +from typing import Callable + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.types.utils import GenericGuardrailAPIInputs + +from .chunked_evaluation import MAX_PROMPT_CHARS, SegmentVerdict +from .credentials import get_metadata +from .exceptions import WonderFenceBlockedError, WonderFenceScanBudgetExceeded + +logger = verbose_proxy_logger.getChild("alice_wonderfence") + +JOINER = "\n" + +# Upper bound on the document ``reconstruct`` will align. Alignment now runs +# per chunk (each <= MAX_PROMPT_CHARS) rather than over the whole document, so +# the cost is O(document / chunk * chunk^2) = O(document * chunk) -- linear in +# document size with the chunk size as the constant, instead of quadratic in the +# document. This bound is a coarse backstop on total alignment work; a MASK on a +# document larger than it fails closed (block). Non-MASK requests of any size are +# unaffected. +RECONSTRUCT_MAX_CHARS = 10 * MAX_PROMPT_CHARS + + +def build_analysis_context( + request_data: dict, + platform: str | None, + context_class: Callable[..., object], +) -> object: + """Build WonderFence AnalysisContext from request data.""" + metadata = get_metadata(request_data) + model_str = request_data.get("model", "") + + provider = None + model_name = model_str + if model_str: + try: + model_name, provider, _, _ = litellm.get_llm_provider(model=model_str) + except Exception: # noqa: BLE001 # best-effort provider parse for telemetry; any failure falls back to manual split and must never break the guardrail + if "/" in model_str: + provider, model_name = model_str.split("/", 1) + + user_id = metadata.get("user_api_key_end_user_id") or metadata.get("end_user_id") or metadata.get("user_id") + + session_id = ( + request_data.get("litellm_session_id") or metadata.get("litellm_session_id") or metadata.get("session_id") + ) + + return context_class( + session_id=session_id, + user_id=user_id, + model_name=model_name, + provider=provider, + platform=platform, + ) + + +def tool_call_arg_segments( + inputs: GenericGuardrailAPIInputs, +) -> tuple[list[int], list[str]]: + """Return (indices, argument strings) for tool calls carrying string args. + + ``inputs["tool_calls"]`` entries are dicts shaped + ``{"function": {"arguments": ""}}``; the argument string is the + caller- or model-controlled payload that reaches the model/client, so it is + scanned. ``indices`` is only used on the response side, where a MASK verdict + is written back in place; on the request side the argument strings are + appended to the joined document as detection-only pieces (see + ``apply_guardrail``). + """ + tool_calls = inputs.get("tool_calls") or [] + indices: list[int] = [] + segments: list[str] = [] + for i, tool_call in enumerate(tool_calls): + fn = tool_call.get("function") if isinstance(tool_call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str) and args.strip(): + indices.append(i) + segments.append(args) + return indices, segments + + +def _description_texts(fn: object) -> list[str]: + """Collect every non-blank ``description`` string under a tool's ``function`` + dict, walking nested JSON-schema parameters so parameter descriptions are + included, not just the top one. + + Iterative (explicit stack) rather than recursive: caller-supplied tool + schemas can nest arbitrarily, and unbounded recursion on request input is a + DoS / stack-overflow risk. Only the strings are returned (no write-back + paths): tool/function descriptions are scanned detection-only, so there is + nothing to mask back into the schema. + """ + out: list[str] = [] + stack: list[object] = [fn] + while stack: + obj = stack.pop() + if isinstance(obj, dict): + for key, value in obj.items(): + if key == "description" and isinstance(value, str) and value.strip(): + out.append(value) + elif isinstance(value, (dict, list)): + stack.append(value) + elif isinstance(obj, list): + stack.extend(item for item in obj if isinstance(item, (dict, list))) + return out + + +def tool_definition_segments(inputs: GenericGuardrailAPIInputs) -> list[str]: + """Return description texts from ``inputs["tools"]`` (detection-only). + + The chat translation layer passes caller-supplied ``inputs["tools"]`` to the + model verbatim, so a tool's ``function.description`` and its nested parameter + descriptions are scanned. Detection-only: they are rendered into the joined + document as extra pieces and can BLOCK/DETECT but are never masked back + (there is no faithful place to splice a redaction into a schema). + """ + tools = inputs.get("tools") or [] + return [ + text + for tool in tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + for text in _description_texts(tool["function"]) + ] + + +def function_definition_segments(request_data: dict) -> list[str]: + """Return description texts from the deprecated top-level ``functions[]``. + + Each entry is shaped like a tool's ``function`` object, so the same + description walker applies. Detection-only, same rationale as + ``tool_definition_segments``. + """ + functions = request_data.get("functions") or [] + return [text for fn in functions if isinstance(fn, dict) for text in _description_texts(fn)] + + +def check_scan_budget( + segments: list[str], + max_scan_chars: int | None, + max_scan_segments: int | None, +) -> None: + """Fail-closed total-work cap; raises ``WonderFenceScanBudgetExceeded`` when + the combined scan characters or segment count exceed the configured limits. + + WonderFence has no batch API (one HTTP POST per string), so without a cap a + single crafted request with thousands of tiny parts or huge content could + amplify into an unbounded number of upstream calls. This check runs before + any WonderFence call and is never subject to ``fail_open`` — a caller must + not be able to bypass scanning by overflowing the cap. + """ + n = len(segments) + if max_scan_segments is not None and n > max_scan_segments: + raise WonderFenceScanBudgetExceeded( + { + "error": f"Alice WonderFence scan budget exceeded: {n} segments > max_scan_segments={max_scan_segments}", + "type": "alice_wonderfence_scan_budget_exceeded", + "limit": "max_scan_segments", + "max_scan_segments": max_scan_segments, + "segments": n, + } + ) + total_chars = sum(len(s) for s in segments) + if max_scan_chars is not None and total_chars > max_scan_chars: + raise WonderFenceScanBudgetExceeded( + { + "error": f"Alice WonderFence scan budget exceeded: {total_chars} chars > max_scan_chars={max_scan_chars}", + "type": "alice_wonderfence_scan_budget_exceeded", + "limit": "max_scan_chars", + "max_scan_chars": max_scan_chars, + "chars": total_chars, + } + ) + + +def _map_index(x: int, ops: Sequence[tuple[str, int, int, int, int]], masked_len: int) -> int | None: + """Map an index in the original joined document to its index in ``masked``. + + Uses the ``SequenceMatcher`` opcodes: an index inside (or at the end of) an + ``equal`` block maps positionally; a boundary that lands at the very start + of a changed block still maps (the range simply begins there); a boundary + that lands *inside* a changed block is ambiguous and returns ``None`` so the + caller fails closed rather than misassigning. + """ + for tag, i1, i2, j1, _j2 in ops: + if i1 <= x < i2 or (x == i2 and tag == "equal"): + if tag == "equal": + return j1 + (x - i1) + return j1 if x == i1 else None + return masked_len + + +_ChunkEntry = tuple[int, int, int, str, Sequence[tuple[str, int, int, int, int]] | None] + + +def _chunk_entries(masked_chunks: list[tuple[str, str]]) -> tuple[list[_ChunkEntry], str]: + """Build the per-chunk position map. + + One entry per owned region: ``(orig_start, orig_end, masked_start, + masked_owned, opcodes|None)`` with cumulative offsets in both original and + masked space. Opcodes are computed only for chunks the service actually + changed (each aligned over <= one chunk, so the alignment cost is bounded by + the chunk size); unchanged chunks map by a fixed offset with no alignment. + Returns the entries and the reassembled masked document. + """ + entries: list[_ChunkEntry] = [] + o_off = 0 + m_off = 0 + for original_owned, masked_owned in masked_chunks: + ops = ( + None + if original_owned == masked_owned + else SequenceMatcher(None, original_owned, masked_owned, autojunk=False).get_opcodes() + ) + entries.append((o_off, o_off + len(original_owned), m_off, masked_owned, ops)) + o_off += len(original_owned) + m_off += len(masked_owned) + return entries, "".join(m for _, m in masked_chunks) + + +def _map_pos(x: int, entries: list[_ChunkEntry], masked_len: int) -> int | None: + """Map original index ``x`` to its masked index via the owning chunk.""" + for o_start, o_end, m_start, masked_owned, ops in entries: + if o_start <= x < o_end: + local = x - o_start + if ops is None: + return m_start + local + r = _map_index(local, ops, len(masked_owned)) + return None if r is None else m_start + r + return masked_len + + +def _joiner_survives(j: int, entries: list[_ChunkEntry]) -> bool: + """Whether the ``JOINER`` at original index ``j`` survives the mask as an + unmodified ``\\n`` (so parts cannot merge).""" + for o_start, o_end, _m_start, masked_owned, ops in entries: + if o_start <= j < o_end: + if ops is None: + return True + local = j - o_start + return any( + tag == "equal" and i1 <= local < i2 and masked_owned[j1 + (local - i1)] == JOINER + for tag, i1, i2, j1, _j2 in ops + ) + return False + + +def reconstruct(parts: list[str], masked_chunks: list[tuple[str, str]] | None) -> list[str] | None: + """Recover per-part masked text from the per-chunk masked owned regions. + + ``parts`` were joined with ``JOINER`` (a plain ``"\\n"``) into the document + that was scanned; ``masked_chunks`` is the list of ``(owned_original, + owned_masked)`` regions that concatenate back to that document and its masked + form. Alignment is done per chunk (see ``_chunk_entries``), so the cost is + bounded by the chunk size rather than quadratic in the whole document; each + part's char range is mapped through the owning chunk's alignment. + + Fails closed (returns ``None``) when the structure is not recoverable: + ``masked_chunks`` is missing; the document exceeds ``RECONSTRUCT_MAX_CHARS``; + the owned regions do not concatenate back to the join (invariant guard); any + ``JOINER`` between parts does not survive as an unmodified ``\\n`` (a mask + spanning a joiner would merge parts); or a part boundary lands inside a + changed block. Returns one masked string per input part, in order; ``[]`` for + no parts. Assumes masking is span substitution that preserves the non-masked + characters; if the service reflows whitespace the joiner-survival check trips + and we fail closed rather than misassign. + """ + if not parts: + return [] + if masked_chunks is None: + return None + + original = JOINER.join(parts) + if len(original) > RECONSTRUCT_MAX_CHARS: + return None + if "".join(o for o, _ in masked_chunks) != original: + return None + + entries, masked_doc = _chunk_entries(masked_chunks) + masked_len = len(masked_doc) + + starts = [0, *accumulate(len(p) + len(JOINER) for p in parts)][: len(parts)] + ranges = [(s, s + len(p)) for s, p in zip(starts, parts)] + joiners = [end for (_s, end) in ranges[:-1]] + + if not all(_joiner_survives(j, entries) for j in joiners): + return None + + mapped = [(_map_pos(s, entries, masked_len), _map_pos(e, entries, masked_len)) for s, e in ranges] + if any(ms is None or me is None or ms > me for ms, me in mapped): + return None + return [masked_doc[ms:me] for ms, me in mapped] + + +def block_detail(blocked: list[SegmentVerdict], guardrail_name: str, block_message: str) -> dict: + detections: list = [] + correlation_ids: list[str] = [] + for v in blocked: + detections.extend(v.detections) + correlation_ids.extend(v.correlation_ids) + detail: dict = { + "error": block_message, + "type": "alice_wonderfence_content_policy_violation", + "guardrail_name": guardrail_name, + "action": "BLOCK", + "wonderfence_correlation_id": correlation_ids[0] if correlation_ids else None, + "wonderfence_correlation_ids": correlation_ids, + } + if detections: + detail["detections"] = [d.model_dump() if hasattr(d, "model_dump") else d for d in detections] + return detail + + +def raise_if_blocked(verdicts: list[SegmentVerdict], guardrail_name: str, block_message: str) -> None: + """Raise ``WonderFenceBlockedError`` if any verdict is BLOCK, aggregating + detections / correlation ids across all blocked verdicts.""" + blocked = [v for v in verdicts if v.action == "BLOCK"] + if blocked: + raise WonderFenceBlockedError(block_detail(blocked, guardrail_name, block_message)) + + +def _masked_value(verdict: SegmentVerdict, guardrail_name: str, label: str) -> str | None: + """Return the replacement string for a MASK verdict (logging as a side + effect), or None for DETECT/NO_ACTION. The caller writes it to the slot the + segment came from.""" + correlation_id = verdict.correlation_ids[0] if verdict.correlation_ids else None + if verdict.action == "MASK": + logger.info( + "Alice WonderFence (apply_guardrail): MASK applied to %s guardrail=%s correlation_id=%s", + label, + guardrail_name, + correlation_id, + ) + return verdict.masked_text if verdict.masked_text is not None else "[MASKED]" + if verdict.action == "DETECT": + logger.warning( + "Alice WonderFence (apply_guardrail): DETECT %s guardrail=%s correlation_id=%s", + label, + guardrail_name, + correlation_id, + ) + return None + + +def apply_response_verdicts( + inputs: GenericGuardrailAPIInputs, + text_verdicts: list[SegmentVerdict], + tool_indices: list[int], + tool_verdicts: list[SegmentVerdict], + guardrail_name: str, + block_message: str, +) -> GenericGuardrailAPIInputs: + """Response-side write-back: index-aligned MASK into ``texts`` (per choice) + and ``tool_calls[i].function.arguments`` (model-generated). + + BLOCK across any segment raises first. Response text is written per-index + (never joined) because the handler's response write-back is purely + positional over the returned ``texts`` list and has no ``structured_messages`` + path, so collapsing choices into one string would dump every choice's text + into choice 0. + """ + raise_if_blocked([*text_verdicts, *tool_verdicts], guardrail_name, block_message) + + texts = inputs.get("texts") or [] + for idx, verdict in enumerate(text_verdicts): + masked = _masked_value(verdict, guardrail_name, "response text") + if masked is not None: + texts[idx] = masked + inputs["texts"] = texts + + tool_calls = inputs.get("tool_calls") or [] + for idx, verdict in zip(tool_indices, tool_verdicts): + masked = _masked_value(verdict, guardrail_name, "tool_call args") + if masked is not None: + tool_calls[idx]["function"]["arguments"] = masked + + return inputs diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c7cdfaad780..bd2875d3e92 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -63,6 +63,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( XecGuardConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.alice_wonderfence import ( + WonderFenceGuardrailConfigModel, +) """ Pydantic object defining how to set guardrails on litellm proxy @@ -134,6 +137,7 @@ class SupportedGuardrailIntegrations(Enum): HEADROOM = "headroom" COMPRESR = "compresr" STRAIKER = "straiker" + ALICE_WONDERFENCE = "alice_wonderfence" class Role(Enum): @@ -999,6 +1003,7 @@ class LitellmParams( QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, SingulrGuardrailConfigModel, + WonderFenceGuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: str | list[str] | Mode = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/alice_wonderfence.py b/litellm/types/proxy/guardrails/guardrail_hooks/alice_wonderfence.py new file mode 100644 index 00000000000..e6655284fd1 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/alice_wonderfence.py @@ -0,0 +1,76 @@ +"""Alice WonderFence guardrail configuration models.""" + +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class WonderFenceGuardrailConfigModel(GuardrailConfigModel): + """Configuration parameters for the Alice WonderFence guardrail. + + Resolution order for ``api_key`` and ``app_id`` (highest first): + API-key metadata → team metadata → request metadata (only when + ``allow_request_metadata_override`` is True) → ``api_key`` falls back to + the configured default below or the ``ALICE_API_KEY`` env var; ``app_id`` + has no default. + + By default, request-body metadata is ignored so a caller cannot bypass + an admin-pinned WonderFence ``app_id`` / ``api_key`` on their virtual + key. Enable ``allow_request_metadata_override`` for trusted-gateway + deployments that legitimately need request-level overrides. + """ + + api_key: Optional[str] = Field( + default=None, + description="Default API key for WonderFence. Overridable via API-key / team metadata, or via request metadata (alice_wonderfence_api_key) only when allow_request_metadata_override is True. Env: ALICE_API_KEY.", + ) + allow_request_metadata_override: Optional[bool] = Field( + default=False, + description="When True, allow alice_wonderfence_api_key and alice_wonderfence_app_id in request metadata as a last-resort source (after API-key and team metadata). Default False so caller-controlled request fields cannot bypass admin-pinned credentials.", + ) + api_base: Optional[str] = Field( + default=None, + description="Override for WonderFence API base URL.", + ) + api_timeout: Optional[float] = Field( + default=10.0, + description="Timeout in seconds for API calls.", + ) + platform: Optional[str] = Field( + default=None, + description="Cloud platform (e.g., aws, azure, databricks).", + ) + fail_open: Optional[bool] = Field( + default=False, + description="When True, proceed with the request/response if WonderFence is unreachable. BLOCK actions are always enforced. Default: False (fail closed).", + ) + block_message: Optional[str] = Field( + default="Content violates our policies and has been blocked", + description="User-facing error message returned when content is blocked.", + ) + debug: Optional[bool] = Field( + default=False, + description="Set guardrail logger to DEBUG level.", + ) + max_cached_clients: Optional[int] = Field( + default=10, + description="Max SDK clients cached per guardrail (LRU, keyed by api_key). Env: ALICE_MAX_CACHED_CLIENTS.", + ) + connection_pool_limit: Optional[int] = Field( + default=None, + description="Max connections per SDK client HTTP pool. Env: ALICE_CONNECTION_POOL_LIMIT.", + ) + max_scan_chars: Optional[int] = Field( + default=1_000_000, + description="Total-work cap (fail-closed DoS backstop): reject a request/response whose combined scan characters (message text plus tool-call args and tool/function descriptions) exceed this before any WonderFence call. Bounds upstream call amplification since WonderFence has no batch API. Env: ALICE_MAX_SCAN_CHARS.", + ) + max_scan_segments: Optional[int] = Field( + default=1_000, + description="Total-work cap (fail-closed DoS backstop): reject a request/response carrying more than this many scan segments (message text parts, tool-call args, tool/function descriptions) before any WonderFence call. Env: ALICE_MAX_SCAN_SEGMENTS.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Alice WonderFence Guardrail" diff --git a/tests/local_testing/test_configs/test_alice_config.yaml b/tests/local_testing/test_configs/test_alice_config.yaml new file mode 100644 index 00000000000..b9ea678e579 --- /dev/null +++ b/tests/local_testing/test_configs/test_alice_config.yaml @@ -0,0 +1,20 @@ +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "alice-wonderfence" + litellm_params: + guardrail: alice_wonderfence + mode: ["during_call", "post_call"] # Test both input and output + api_key: os.environ/ALICE_API_KEY + allow_request_metadata_override: true # app_id supplied per-request via metadata.alice_wonderfence_app_id + api_timeout: 20.0 # Timeout in seconds (default: 20.0) + platform: aws # Optional: Cloud platform (aws, azure, databricks, etc.) + default_on: true + + +litellm_settings: + set_verbose: true \ No newline at end of file diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/conftest.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/conftest.py new file mode 100644 index 00000000000..ced0a4d3de2 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/conftest.py @@ -0,0 +1,133 @@ +"""Shared fixtures for Alice WonderFence guardrail tests.""" + +import sys +from unittest.mock import AsyncMock, Mock + +import pytest + + +def _install_sdk_stub(monkeypatch, client_factory=None): + """Install a stub ``wonderfence_sdk`` module so the guardrail can import it.""" + sdk = Mock() + client_pkg = Mock() + models_pkg = Mock() + + factory = client_factory or (lambda **kwargs: Mock(close=AsyncMock())) + client_pkg.WonderFenceV2Client = Mock(side_effect=factory) + sdk.client = client_pkg + + models_pkg.AnalysisContext = Mock(return_value=Mock()) + sdk.models = models_pkg + + monkeypatch.setitem(sys.modules, "wonderfence_sdk", sdk) + monkeypatch.setitem(sys.modules, "wonderfence_sdk.client", client_pkg) + monkeypatch.setitem(sys.modules, "wonderfence_sdk.models", models_pkg) + return sdk + + +def _make_guardrail(monkeypatch, **overrides): + """Build a WonderFenceGuardrail with stubbed SDK and a mock V2 client.""" + from litellm.types.guardrails import GuardrailEventHooks + + mock_client = Mock() + mock_client.evaluate_prompt = AsyncMock() + mock_client.evaluate_response = AsyncMock() + mock_client.close = AsyncMock() + + _install_sdk_stub(monkeypatch, client_factory=lambda **kwargs: mock_client) + + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.alice_wonderfence import ( + WonderFenceGuardrail, + ) + + kwargs = dict( + guardrail_name="wonderfence-test", + api_key="default-api-key", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + default_on=True, + ) + kwargs.update(overrides) + guardrail = WonderFenceGuardrail(**kwargs) + return guardrail, mock_client + + +def _request_data(**overrides): + """Build a request-data dict. + + Default metadata pins ``alice_wonderfence_app_id`` on + ``user_api_key_metadata`` (admin-controlled) so the request resolves + cleanly under the safe-by-default precedence model. Tests that want to + drive the value through request metadata must (a) construct a guardrail + with ``allow_request_metadata_override=True`` and (b) pass the value via + the ``metadata`` kwarg explicitly. + """ + metadata = overrides.pop("metadata", None) + if metadata is None: + metadata = {"user_api_key_metadata": {"alice_wonderfence_app_id": "test-app"}} + base = {"model": "gpt-4", "metadata": metadata} + base.update(overrides) + return base + + +def _make_logging_obj(): + """Build a real ``LiteLLMLoggingObj``. + + The post_call bridge stashes resolved credentials on a private attribute of + this object, so tests must use the real class (not a Mock, whose attribute + auto-creation would mask whether the attribute is genuinely settable and + readable) to validate that the stash survives request -> response. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=0, + litellm_call_id="alice-wonderfence-test", + function_id="alice-wonderfence-test", + ) + + +@pytest.fixture +def guardrail_and_client(monkeypatch): + g, c = _make_guardrail(monkeypatch) + # Pre-seed cache so apply_guardrail uses our mock without rebuilding. + g._client_cache["default-api-key"] = c + return g, c + + +@pytest.fixture +def install_sdk_stub(monkeypatch): + """Expose ``_install_sdk_stub`` as a fixture for tests that need direct access.""" + + def _factory(client_factory=None): + return _install_sdk_stub(monkeypatch, client_factory=client_factory) + + return _factory + + +@pytest.fixture +def make_guardrail(monkeypatch): + """Expose ``_make_guardrail`` as a fixture.""" + + def _factory(**overrides): + return _make_guardrail(monkeypatch, **overrides) + + return _factory + + +@pytest.fixture +def make_request_data(): + """Expose ``_request_data`` as a fixture.""" + return _request_data + + +@pytest.fixture +def make_logging_obj(): + """Expose ``_make_logging_obj`` as a fixture.""" + return _make_logging_obj diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail.py new file mode 100644 index 00000000000..d9c737dc980 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail.py @@ -0,0 +1,534 @@ +"""Tests for ``apply_guardrail`` text-side BLOCK/MASK/DETECT/NO_ACTION + core scanning path. + +Tool-call / tool-definition / legacy functions[] scanning lives in +``test_apply_guardrail_tools.py``; fail-open/closed, missing secrets, and helpers +live in ``test_apply_guardrail_failmodes.py``. +""" + +from unittest.mock import Mock + +import pytest +from fastapi import HTTPException + +# ----------------------------- BLOCK ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_block_action(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "BLOCK" + detection = Mock() + detection.model_dump = Mock(return_value={"policy_name": "x", "confidence": 0.9}) + result_obj.detections = [detection] + result_obj.correlation_id = "corr-1" + client.evaluate_prompt.return_value = result_obj + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert exc.value.detail["action"] == "BLOCK" + assert exc.value.detail["wonderfence_correlation_id"] == "corr-1" + assert exc.value.detail["error"] == ("Content violates our policies and has been blocked") + assert exc.value.detail["detections"][0]["policy_name"] == "x" + + +@pytest.mark.asyncio +async def test_apply_guardrail_block_uses_custom_block_message(make_guardrail, make_request_data): + guardrail, client = make_guardrail(block_message="custom blocked text") + guardrail._client_cache["default-api-key"] = client + result_obj = Mock() + result_obj.action = "BLOCK" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.detail["error"] == "custom blocked text" + + +@pytest.mark.asyncio +async def test_block_not_bypassed_by_fail_open(make_guardrail, make_request_data): + guardrail, client = make_guardrail(fail_open=True) + guardrail._client_cache["default-api-key"] = client + result_obj = Mock() + result_obj.action = "BLOCK" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["bad"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + + +# ----------------------------- MASK ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_replaces_scanned_text(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "MASK" + result_obj.action_text = "[REDACTED]" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + out = await guardrail.apply_guardrail( + inputs={"texts": ["sensitive"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["[REDACTED]"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_reconstructs_only_the_flagged_message_part(guardrail_and_client, make_request_data): + """Request side joins the message parts into one document, scans once, and on + MASK reconstructs per-part masked text by aligning the join against the + masked document. Only the flagged part changes; the others survive.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + # One joined call: mask just the sensitive part inside the joined doc. + r.action = "MASK" + r.action_text = prompt.replace("sensitive content", "[REDACTED]") + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + out = await guardrail.apply_guardrail( + inputs={"texts": ["first", "ack", "sensitive content"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["first", "ack", "[REDACTED]"] + client.evaluate_prompt.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_apply_guardrail_scans_non_user_role_segments(guardrail_and_client, make_request_data): + """Bypass regression: blocked content in a system/assistant/tool message + must still BLOCK. The translation layer already strips system/tool when the + guardrail is configured to skip them, so whatever remains in ``texts`` is + scanned regardless of role; the hook must not re-filter to user-only.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "disallowed system instruction" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "structured_messages": [ + {"role": "system", "content": "disallowed system instruction"}, + {"role": "user", "content": "hello"}, + ], + "texts": ["disallowed system instruction", "hello"], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert exc.value.detail["action"] == "BLOCK" + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_replaces_scanned_text_response(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "MASK" + result_obj.action_text = "[REDACTED]" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_response.return_value = result_obj + + out = await guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data=make_request_data(), + input_type="response", + ) + assert out["texts"] == ["[REDACTED]"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_fallback_when_action_text_is_none(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "MASK" + result_obj.action_text = None + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + out = await guardrail.apply_guardrail( + inputs={"texts": ["a"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["[MASKED]"] + + +# ----------------------------- DETECT / NO_ACTION ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_action_passthrough(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "NO_ACTION" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + out = await guardrail.apply_guardrail( + inputs={"texts": ["safe"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["safe"] + client.evaluate_prompt.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_apply_guardrail_detect_action_passes_through(guardrail_and_client, make_request_data): + """DETECT action logs a warning but does not block or mutate inputs.""" + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "DETECT" + result_obj.detections = [] + result_obj.correlation_id = "corr-detect" + client.evaluate_prompt.return_value = result_obj + + out = await guardrail.apply_guardrail( + inputs={"texts": ["watch me"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["watch me"] + client.evaluate_prompt.assert_awaited_once() + + +# ----------------------------- core path / app_id passthrough ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_passes_app_id_per_call(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "NO_ACTION" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={"user_api_key_metadata": {"alice_wonderfence_app_id": "tenant-A"}}), + input_type="request", + ) + kwargs = client.evaluate_prompt.call_args.kwargs + assert kwargs["app_id"] == "tenant-A" + assert kwargs["prompt"] == "hi" + assert kwargs["custom_fields"] is None + + +@pytest.mark.asyncio +async def test_apply_guardrail_response_path_passes_app_id(make_guardrail, make_request_data): + guardrail, client = make_guardrail() + guardrail._client_cache["default-api-key"] = client + result_obj = Mock() + result_obj.action = "NO_ACTION" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_response.return_value = result_obj + + await guardrail.apply_guardrail( + inputs={"texts": ["resp"]}, + request_data=make_request_data(metadata={"user_api_key_metadata": {"alice_wonderfence_app_id": "tenant-B"}}), + input_type="response", + ) + kwargs = client.evaluate_response.call_args.kwargs + assert kwargs["app_id"] == "tenant-B" + assert kwargs["response"] == "resp" + + +@pytest.mark.asyncio +async def test_apply_guardrail_joins_all_message_parts_into_one_call(guardrail_and_client, make_request_data): + """Every message part is scanned, but as a single joined document in ONE + Alice call (call volume scales with size, not message count). The join uses + a plain newline so cross-part content is seen whole.""" + guardrail, client = guardrail_and_client + result_obj = Mock() + result_obj.action = "NO_ACTION" + result_obj.detections = [] + result_obj.correlation_id = None + client.evaluate_prompt.return_value = result_obj + + await guardrail.apply_guardrail( + inputs={"texts": ["t1", "t2", "t3"]}, + request_data=make_request_data(), + input_type="request", + ) + client.evaluate_prompt.assert_awaited_once() + assert client.evaluate_prompt.call_args.kwargs["prompt"] == "t1\nt2\nt3" + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_earlier_user_turn(guardrail_and_client, make_request_data): + """Bypass regression: disallowed content in an earlier user turn followed by + a benign final turn must still BLOCK. The old last-only path only saw the + benign final message and let the request through.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "disallowed" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "structured_messages": [ + {"role": "user", "content": "disallowed"}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "hello"}, + ], + "texts": ["disallowed", "ok", "hello"], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert exc.value.detail["action"] == "BLOCK" + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_when_oversized_message_trips_in_late_chunk( + guardrail_and_client, make_request_data +): + """A single user message over the prompt limit is chunked; a BLOCK in a + non-first chunk still blocks the request.""" + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence import ( + chunked_evaluation, + ) + + guardrail, client = guardrail_and_client + long_prompt = ("safe " * 5000) + "TRIPWIRE" + assert len(long_prompt) > chunked_evaluation.MAX_PROMPT_CHARS + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "TRIPWIRE" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": [long_prompt]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert client.evaluate_prompt.call_count > 1 + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_text_short_circuits(guardrail_and_client, make_request_data): + """Empty inputs must skip the SDK call and return inputs unchanged.""" + guardrail, client = guardrail_and_client + out = await guardrail.apply_guardrail( + inputs={"texts": []}, + request_data=make_request_data(), + input_type="request", + ) + assert out == {"texts": []} + client.evaluate_prompt.assert_not_awaited() + client.evaluate_response.assert_not_awaited() + + +# ----------------------------- join: cross-part visibility ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_phrase_split_across_message_parts(guardrail_and_client, make_request_data): + """Two content parts that individually look benign are joined into one + document, so a phrase split across the part boundary is seen in a single + scan and still BLOCKs (the join replaces the old cross-segment windows).""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + # Neither part alone contains the whole phrase; the joined document does. + r.action = "BLOCK" if ("make a b" in prompt and "omb" in prompt) else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["how to make a b", "omb please"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + client.evaluate_prompt.assert_awaited_once() + + +# ----------------------------- join: MASK reconstruction write-back ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_writes_back_to_the_correct_message(guardrail_and_client, make_request_data): + """A real PII MASK on the joined document reconstructs per-part masked text + and writes it back to the message that carried the PII, leaving the others + intact.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "MASK" + r.action_text = prompt.replace("john@example.com", "[EMAIL]") + r.detections = [] + r.correlation_id = "corr-mask" + return r + + client.evaluate_prompt.side_effect = evaluate + + out = await guardrail.apply_guardrail( + inputs={"texts": ["hello there", "my email is john@example.com", "thanks"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["hello there", "my email is [EMAIL]", "thanks"] + client.evaluate_prompt.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_reconstruction_failure_fails_closed(guardrail_and_client, make_request_data): + """If masking destroys a joiner (parts would merge), reconstruction cannot + safely attribute the redaction, so the request is blocked rather than + silently misassigned or passed through unmasked.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "MASK" + r.action_text = prompt.replace("\n", "") # destroys the joiner -> parts merge + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["alpha", "beta"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + + +# ----------------------------- total-work cap ----------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_rejects_over_segment_cap_without_scanning(make_guardrail, make_request_data): + guardrail, client = make_guardrail(max_scan_segments=3) + guardrail._client_cache["default-api-key"] = client + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["a", "b", "c", "d"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert exc.value.detail["limit"] == "max_scan_segments" + client.evaluate_prompt.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_apply_guardrail_cap_is_not_bypassed_by_fail_open(make_guardrail, make_request_data): + """The cap is a config/abuse guard, never fail-open: an oversized request + is rejected 400 even with fail_open=True and the SDK is never called.""" + guardrail, client = make_guardrail(max_scan_chars=10, fail_open=True) + guardrail._client_cache["default-api-key"] = client + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["x" * 50]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert exc.value.detail["limit"] == "max_scan_chars" + client.evaluate_prompt.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_apply_guardrail_mask_on_oversized_document_fails_closed(make_guardrail, make_request_data): + """A MASK on a document too large to reconstruct within the bounded + alignment cost fails closed (block) rather than run the quadratic + SequenceMatcher on the event loop or forward unmasked content.""" + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.processing import ( + RECONSTRUCT_MAX_CHARS, + ) + + # Keep the char cap high enough to reach scanning, but exceed the + # reconstruction bound so MASK cannot be applied. + guardrail, client = make_guardrail(max_scan_chars=RECONSTRUCT_MAX_CHARS * 2) + guardrail._client_cache["default-api-key"] = client + big = "a " * RECONSTRUCT_MAX_CHARS # > RECONSTRUCT_MAX_CHARS chars + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "MASK" + r.action_text = "[REDACTED]" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": [big]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail_failmodes.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail_failmodes.py new file mode 100644 index 00000000000..7863b029626 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail_failmodes.py @@ -0,0 +1,184 @@ +"""Tests for ``apply_guardrail`` fail-open/fail-closed behavior, missing secrets, and helpers.""" + +import sys +from unittest.mock import Mock + +import pytest +from fastapi import HTTPException + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_app_id_fail_closed_returns_500(guardrail_and_client, make_request_data): + """Missing app_id follows the fail_open pattern: fail_open=False → HTTP 500.""" + guardrail, _ = guardrail_and_client + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={}), + input_type="request", + ) + assert exc.value.status_code == 500 + assert "Error in Alice WonderFence Guardrail" in exc.value.detail["error"] + assert "alice_wonderfence_app_id" in exc.value.detail["exception"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_api_key_fail_closed_returns_500(monkeypatch, make_guardrail, make_request_data): + """Missing api_key follows the fail_open pattern: fail_open=False → HTTP 500.""" + monkeypatch.delenv("ALICE_API_KEY", raising=False) + guardrail, _ = make_guardrail(api_key=None) + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 500 + assert "Error in Alice WonderFence Guardrail" in exc.value.detail["error"] + assert "alice_wonderfence_api_key" in exc.value.detail["exception"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_app_id_fail_open_returns_500(make_guardrail, make_request_data): + """Missing app_id is a config error: never fail-open, even with fail_open=True.""" + guardrail, _ = make_guardrail(fail_open=True) + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={}), + input_type="request", + ) + assert exc.value.status_code == 500 + assert "alice_wonderfence_app_id" in exc.value.detail["exception"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_api_key_fail_open_returns_500(monkeypatch, make_guardrail, make_request_data): + """Missing api_key is a config error: never fail-open, even with fail_open=True.""" + monkeypatch.delenv("ALICE_API_KEY", raising=False) + guardrail, _ = make_guardrail(api_key=None, fail_open=True) + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 500 + assert "alice_wonderfence_api_key" in exc.value.detail["exception"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_fail_open_swallows_transport_error(make_guardrail, make_request_data): + guardrail, client = make_guardrail(fail_open=True) + guardrail._client_cache["default-api-key"] = client + client.evaluate_prompt.side_effect = RuntimeError("network down") + + inputs = {"texts": ["original"]} + out = await guardrail.apply_guardrail( + inputs=inputs, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["original"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_fail_closed_returns_500(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + client.evaluate_prompt.side_effect = RuntimeError("network down") + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 500 + assert "Error in Alice WonderFence Guardrail" in exc.value.detail["error"] + + +@pytest.mark.asyncio +async def test_client_error_4xx_not_fail_open(make_guardrail, make_request_data): + """A persistent client/config error (the SDK raises Exception("<4xx>:...") + for an invalid/revoked api_key or app_id) must NOT fail open, even with + fail_open=True: it would silently disable scanning for every affected + request. It surfaces as HTTP 500 and the original inputs are not returned.""" + guardrail, client = make_guardrail(fail_open=True) + guardrail._client_cache["default-api-key"] = client + client.evaluate_prompt.side_effect = Exception('401:{"detail":"invalid api key"}') + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["original"]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 500 + client.evaluate_prompt.assert_awaited() + + +@pytest.mark.asyncio +async def test_transient_5xx_and_429_still_fail_open(make_guardrail, make_request_data): + """5xx and 429 are transient (matching the SDK's retry classification), so + fail_open still lets the request through unscanned.""" + for status in ("503", "429"): + guardrail, client = make_guardrail(fail_open=True) + guardrail._client_cache["default-api-key"] = client + client.evaluate_prompt.side_effect = Exception(f"{status}:upstream unavailable") + + out = await guardrail.apply_guardrail( + inputs={"texts": ["original"]}, + request_data=make_request_data(), + input_type="request", + ) + assert out["texts"] == ["original"], f"status {status} should fail open" + + +@pytest.mark.asyncio +async def test_malformed_override_does_not_fail_open(make_guardrail, make_request_data): + """A non-string request-metadata app_id override must not slip through under + fail_open: it resolves to a config error (500), not a swallowed exception + that skips scanning. The SDK is never called with a malformed value.""" + guardrail, client = make_guardrail(fail_open=True, allow_request_metadata_override=True) + guardrail._client_cache["default-api-key"] = client + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={"alice_wonderfence_app_id": ["not", "a", "string"]}), + input_type="request", + ) + assert exc.value.status_code == 500 + assert "alice_wonderfence_app_id" in exc.value.detail["exception"] + client.evaluate_prompt.assert_not_awaited() + + +def test_get_config_model(make_guardrail): + from litellm.types.proxy.guardrails.guardrail_hooks.alice_wonderfence import ( + WonderFenceGuardrailConfigModel, + ) + + guardrail, _ = make_guardrail() + assert guardrail.get_config_model() is WonderFenceGuardrailConfigModel + + +def test_build_analysis_context_falls_back_to_slash_split(monkeypatch, make_guardrail): + """When ``litellm.get_llm_provider`` raises, fall back to ``provider/model`` split.""" + import litellm + + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.processing import ( + build_analysis_context, + ) + + guardrail, _ = make_guardrail() + + def boom(model): + raise ValueError("unknown provider") + + monkeypatch.setattr(litellm, "get_llm_provider", boom) + build_analysis_context({"model": "myorg/custom-llm"}, guardrail.platform, guardrail._AnalysisContext) + + AnalysisContext = sys.modules["wonderfence_sdk.models"].AnalysisContext + kwargs = AnalysisContext.call_args.kwargs + assert kwargs["provider"] == "myorg" + assert kwargs["model_name"] == "custom-llm" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail_tools.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail_tools.py new file mode 100644 index 00000000000..8bfb873e9f7 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_apply_guardrail_tools.py @@ -0,0 +1,393 @@ +"""Tests for ``apply_guardrail`` scanning of tool calls, tool definitions, and legacy functions[].""" + +from unittest.mock import Mock + +import pytest +from fastapi import HTTPException + + +def _tool_call(arguments, name="send_email"): + return { + "id": "call_1", + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_tool_call_arguments(guardrail_and_client, make_request_data): + """Bypass regression: blocked content in tool_calls[].function.arguments must + BLOCK. tool_calls reach the model but were never scanned (texts-only).""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "texts": ["please run the tool"], + "tool_calls": [_tool_call('{"body": "DISALLOWED payload"}')], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + assert exc.value.detail["action"] == "BLOCK" + + +@pytest.mark.asyncio +async def test_apply_guardrail_request_tool_call_args_mask_fails_closed(guardrail_and_client, make_request_data): + """On the request side, tool-call args are detection-only pieces in the join. + A MASK that redacts an arg cannot be spliced back into the wire-format + arguments string, so forwarding the original unredacted value would leak it; + the request fails closed (block) instead.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "MASK" + r.action_text = prompt.replace("secret value", "[REDACTED]") + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "texts": ["benign"], + "tool_calls": [_tool_call('{"body": "secret value"}')], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_detect_on_tool_call_args_passes_through(guardrail_and_client, make_request_data): + """A DETECT verdict on a tool-call argument logs but does not block or mutate + the arguments (symmetric with the text-side DETECT behavior).""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "DETECT" if "watch me" in prompt else "NO_ACTION" + r.action_text = None + r.detections = [] + r.correlation_id = "corr-detect" + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = {"texts": ["benign"], "tool_calls": [_tool_call('{"x": "watch me"}')]} + out = await guardrail.apply_guardrail( + inputs=inputs, + request_data=make_request_data(), + input_type="request", + ) + assert out["tool_calls"][0]["function"]["arguments"] == '{"x": "watch me"}' + assert out["texts"] == ["benign"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_scans_tool_calls_when_no_texts(guardrail_and_client, make_request_data): + """An assistant message can carry tool_calls with no text content, so texts + is empty; the hook must still scan the tool-call arguments (the old + empty-texts early return skipped them).""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [_tool_call('{"x": "DISALLOWED"}')]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_response_tool_call_arguments(guardrail_and_client, make_request_data): + """Model-generated tool-call arguments on the response side are scanned too.""" + guardrail, client = guardrail_and_client + + def evaluate(response, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in response else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_response.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [_tool_call('{"x": "DISALLOWED"}')]}, + request_data=make_request_data(), + input_type="response", + ) + assert exc.value.status_code == 400 + + +def _tool_def(description="a helpful tool", param_desc=None): + fn = { + "name": "do_thing", + "description": description, + "parameters": {"type": "object", "properties": {}}, + } + if param_desc is not None: + fn["parameters"]["properties"]["city"] = { + "type": "string", + "description": param_desc, + } + return {"type": "function", "function": fn} + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_tool_definition_description(guardrail_and_client, make_request_data): + """Blocked content in tools[].function.description must BLOCK; tool defs are + forwarded to the model but were previously unscanned.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "texts": ["use the tool"], + "tools": [_tool_def(description="DISALLOWED instructions here")], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail(inputs=inputs, request_data=make_request_data(), input_type="request") + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_tool_parameter_description(guardrail_and_client, make_request_data): + """Nested parameter descriptions are scanned too, not just the top-level one.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "texts": ["hi"], + "tools": [_tool_def(description="benign", param_desc="DISALLOWED payload")], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail(inputs=inputs, request_data=make_request_data(), input_type="request") + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_tool_definition_mask_fails_closed(guardrail_and_client, make_request_data): + """Tool definitions are detection-only; a MASK that would redact a + description cannot be spliced back into the schema, so the request fails + closed rather than forward the original unredacted description.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "MASK" + r.action_text = prompt.replace("secret", "[REDACTED]") + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + inputs = { + "texts": ["hi"], + "tools": [_tool_def(description="contains secret stuff")], + } + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail(inputs=inputs, request_data=make_request_data(), input_type="request") + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_scans_tools_when_no_texts_or_tool_calls(guardrail_and_client, make_request_data): + """A request carrying only tool definitions must still be scanned.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": [], "tools": [_tool_def(description="DISALLOWED")]}, + request_data=make_request_data(), + input_type="request", + ) + assert exc.value.status_code == 400 + + +def _legacy_function(description="a function", param_desc=None): + fn = { + "name": "do_thing", + "description": description, + "parameters": {"type": "object", "properties": {}}, + } + if param_desc is not None: + fn["parameters"]["properties"]["city"] = { + "type": "string", + "description": param_desc, + } + return fn + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_legacy_function_description(guardrail_and_client, make_request_data): + """Blocked content in the deprecated functions[].description (read from + request_data, not inputs) must BLOCK.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(functions=[_legacy_function(description="DISALLOWED instructions")]), + input_type="request", + ) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_on_legacy_function_parameter_description(guardrail_and_client, make_request_data): + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(functions=[_legacy_function(description="ok", param_desc="DISALLOWED")]), + input_type="request", + ) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_scans_legacy_functions_when_no_other_content(guardrail_and_client, make_request_data): + """A request whose only scannable content is functions[] is still scanned.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "BLOCK" if "DISALLOWED" in prompt else "NO_ACTION" + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": []}, + request_data=make_request_data(functions=[_legacy_function(description="DISALLOWED")]), + input_type="request", + ) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_legacy_function_detect_does_not_mutate(guardrail_and_client, make_request_data): + """A DETECT verdict on a function definition logs but does not rewrite it.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "DETECT" if "watch" in prompt else "NO_ACTION" + r.action_text = None + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + request_data = make_request_data(functions=[_legacy_function(description="watch this")]) + out = await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + assert out is not None + assert request_data["functions"][0]["description"] == "watch this" + + +@pytest.mark.asyncio +async def test_apply_guardrail_legacy_function_definition_mask_fails_closed(guardrail_and_client, make_request_data): + """Legacy functions[] descriptions are detection-only; a MASK that would + redact one fails closed rather than forward the original unredacted value.""" + guardrail, client = guardrail_and_client + + def evaluate(prompt, **kwargs): + r = Mock() + r.action = "MASK" + r.action_text = prompt.replace("secret", "[REDACTED]") + r.detections = [] + r.correlation_id = None + return r + + client.evaluate_prompt.side_effect = evaluate + + request_data = make_request_data(functions=[_legacy_function(description="contains secret stuff")]) + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + assert exc.value.status_code == 400 + assert request_data["functions"][0]["description"] == "contains secret stuff" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_chunked_evaluation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_chunked_evaluation.py new file mode 100644 index 00000000000..8f978330515 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_chunked_evaluation.py @@ -0,0 +1,252 @@ +"""Tests for the WonderFence-agnostic chunk + parallel-evaluate + aggregate unit.""" + +import asyncio +from unittest.mock import Mock + +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.chunked_evaluation import ( + MAX_PROMPT_CHARS, + SegmentVerdict, + _split_text, + evaluate_segments, +) + + +def _result(action, action_text=None, detections=None, correlation_id=None): + r = Mock() + r.action = action + r.action_text = action_text + r.detections = detections or [] + r.correlation_id = correlation_id + return r + + +# ----------------------------- _split_text ----------------------------- + + +def test_split_short_text_is_single_chunk(): + assert _split_text("hello world", 10000) == ["hello world"] + + +def test_split_long_text_is_lossless(): + text = " ".join(f"word{i}" for i in range(5000)) + chunks = _split_text(text, 100) + assert len(chunks) > 1 + assert all(len(c) <= 100 for c in chunks) + assert "".join(chunks) == text + + +def test_split_force_splits_oversized_single_token(): + text = "x" * 250 + chunks = _split_text(text, 100) + assert all(len(c) <= 100 for c in chunks) + assert "".join(chunks) == text + + +# ----------------------------- evaluate_segments alignment ----------------------------- + + +@pytest.mark.asyncio +async def test_verdicts_align_one_to_one_with_segments(): + actions = {"a": "BLOCK", "b": "MASK", "c": ""} + + async def evaluate(text): + return _result(actions[text], action_text="[M]" if actions[text] == "MASK" else None) + + verdicts = await evaluate_segments(["a", "b", "c"], evaluate) + assert [v.action for v in verdicts] == ["BLOCK", "MASK", ""] + assert isinstance(verdicts[0], SegmentVerdict) + + +@pytest.mark.asyncio +async def test_mask_verdict_carries_masked_text(): + async def evaluate(text): + return _result("MASK", action_text="[REDACTED]") + + verdicts = await evaluate_segments(["secret"], evaluate) + assert verdicts[0].action == "MASK" + assert verdicts[0].masked_text == "[REDACTED]" + + +# ----------------------------- chunking precedence ----------------------------- + + +@pytest.mark.asyncio +async def test_block_in_non_first_chunk_blocks_whole_segment(): + """A segment split into chunks where only a later chunk trips BLOCK must + still produce a BLOCK verdict; the old last-only path never saw earlier text.""" + segment = ("safe " * 30) + "TRIPWIRE" + + async def evaluate(text): + return _result("BLOCK" if "TRIPWIRE" in text else "") + + verdicts = await evaluate_segments([segment], evaluate, max_chars=50) + assert verdicts[0].action == "BLOCK" + + +@pytest.mark.asyncio +async def test_mask_rejoins_masked_owned_regions_into_full_segment(): + """A multi-chunk segment where the service redacts a token in one chunk (a + real span substitution that preserves surrounding bytes) rejoins into the + fully masked segment. overlap=0 keeps the chunks disjoint for a clean check; + the per-chunk (original, masked) pairs are carried on the verdict.""" + segment = " ".join(f"w{i}" for i in range(40)) + " SECRET " + " ".join(f"v{i}" for i in range(40)) + + async def evaluate(text): + return _result("MASK", action_text=text.replace("SECRET", "[X]")) if "SECRET" in text else _result("") + + chunks = _split_text(segment, 20) + assert len(chunks) > 1 + verdicts = await evaluate_segments([segment], evaluate, max_chars=20, overlap=0) + assert verdicts[0].action == "MASK" + assert verdicts[0].masked_text == segment.replace("SECRET", "[X]") + assert verdicts[0].masked_chunks is not None + assert "".join(o for o, _ in verdicts[0].masked_chunks) == segment + + +@pytest.mark.asyncio +async def test_unmasked_chunks_keep_original_text_on_rejoin(): + """Chunks the service did not mask contribute their original owned text + verbatim; only the masked chunk changes.""" + segment = " ".join(f"w{i}" for i in range(40)) + + async def evaluate(text): + return _result("MASK", action_text=text.replace("w0", "[X]")) if "w0 " in text else _result("") + + verdicts = await evaluate_segments([segment], evaluate, max_chars=20, overlap=0) + assert verdicts[0].action == "MASK" + assert verdicts[0].masked_text == segment.replace("w0", "[X]", 1) + + +@pytest.mark.asyncio +async def test_block_beats_mask_within_segment(): + chunks_seen = [] + + async def evaluate(text): + chunks_seen.append(text) + return _result("BLOCK" if "B" in text else "MASK", action_text="[m]") + + segment = "aaa B" + verdicts = await evaluate_segments([segment], evaluate, max_chars=2) + assert verdicts[0].action == "BLOCK" + + +# ----------------------------- aggregation of detections/correlation ids ----------------------------- + + +@pytest.mark.asyncio +async def test_block_verdict_aggregates_detections_and_correlation_ids(): + d1, d2 = Mock(), Mock() + + async def evaluate(text): + if "x" in text: + return _result("BLOCK", detections=[d1], correlation_id="c1") + return _result("BLOCK", detections=[d2], correlation_id="c2") + + verdicts = await evaluate_segments(["x", "y"], evaluate) + assert verdicts[0].detections == [d1] + assert verdicts[0].correlation_ids == ["c1"] + assert verdicts[1].correlation_ids == ["c2"] + + +# ----------------------------- concurrency cap ----------------------------- + + +@pytest.mark.asyncio +async def test_evaluations_run_in_parallel_under_a_cap(): + state = {"current": 0, "max_seen": 0} + + async def evaluate(text): + state["current"] += 1 + state["max_seen"] = max(state["max_seen"], state["current"]) + await asyncio.sleep(0.01) + state["current"] -= 1 + return _result("") + + segments = [f"s{i}" for i in range(12)] + await evaluate_segments(segments, evaluate, max_concurrency=3) + assert state["max_seen"] > 1, "evaluations did not run concurrently" + assert state["max_seen"] <= 3, "concurrency cap exceeded" + + +def test_max_prompt_chars_is_positive(): + assert isinstance(MAX_PROMPT_CHARS, int) and MAX_PROMPT_CHARS > 0 + + +# ----------------------------- seam overlap (detection / masking across chunk splits) ----------------------------- + + +@pytest.mark.asyncio +async def test_block_phrase_split_across_chunk_seam_is_detected(): + """A blocked phrase straddling an owned-region seam is caught because the + next chunk carries an overlap prefix from the previous owned region, so one + scan sees the phrase whole even though neither disjoint owned region does.""" + segment = "aaaaa BLOCK ME zzzzz" + chunks = _split_text(segment, 12) + assert len(chunks) > 1 + assert all("BLOCK ME" not in c for c in chunks) + + async def evaluate(text): + return _result("BLOCK" if "BLOCK ME" in text else "") + + verdicts = await evaluate_segments([segment], evaluate, max_chars=12, overlap=6) + assert verdicts[0].action == "BLOCK" + + +@pytest.mark.asyncio +async def test_no_overlap_lets_seam_phrase_evade(): + """Control: with overlap disabled there is no prefix, so the same straddling + phrase is seen by neither owned region -- demonstrating what the overlap closes.""" + segment = "aaaaa BLOCK ME zzzzz" + + async def evaluate(text): + return _result("BLOCK" if "BLOCK ME" in text else "") + + verdicts = await evaluate_segments([segment], evaluate, max_chars=12, overlap=0) + assert verdicts[0].action == "" + + +@pytest.mark.asyncio +async def test_single_chunk_segment_evaluates_once(): + calls = [] + + async def evaluate(text): + calls.append(text) + return _result("") + + await evaluate_segments(["short benign text"], evaluate, max_chars=10000) + assert calls == ["short benign text"] + + +@pytest.mark.asyncio +async def test_mask_straddling_a_seam_fails_closed_as_block(): + """A MASK whose redaction reaches into a chunk's overlap prefix (content + straddling, or within `overlap` of, an owned-region seam) cannot be stitched + without double-counting the overlap, so it fails closed as BLOCK rather than + leak the un-redacted half. This is the preempt for the seam-mask leak.""" + segment = "aaaaa SECRET HERE zzzzz" + chunks = _split_text(segment, 12) + assert len(chunks) > 1 + + async def evaluate(text): + # The chunk that sees "SECRET HERE" whole (via its overlap prefix) masks + # it; the redaction lands in the prefix bytes -> fail closed. + return _result("MASK", action_text=text.replace("SECRET HERE", "[X]")) if "SECRET HERE" in text else _result("") + + verdicts = await evaluate_segments([segment], evaluate, max_chars=12, overlap=6) + assert verdicts[0].action == "BLOCK" + + +@pytest.mark.asyncio +async def test_independent_segments_are_not_concatenated(): + """Segments are scanned independently (no cross-segment window): a phrase + split across two segments is NOT joined. On the request side, message parts + are joined into one document *before* reaching here; the response side has + independent choices that the model never concatenates.""" + + async def evaluate(text): + return _result("BLOCK" if "BLOCKME" in text else "") + + verdicts = await evaluate_segments(["BLOCK", "ME"], evaluate) + assert [v.action for v in verdicts] == ["", ""] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_client_cache.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_client_cache.py new file mode 100644 index 00000000000..890d7208e57 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_client_cache.py @@ -0,0 +1,195 @@ +"""Tests for the LRU client cache + SDK loader + guardrail initializer.""" + +import sys +from unittest.mock import AsyncMock, Mock + +import pytest + +# ----------------------------- LRU cache ----------------------------- + + +@pytest.mark.asyncio +async def test_get_client_caches_per_api_key(install_sdk_stub): + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.alice_wonderfence import ( + WonderFenceGuardrail, + ) + from litellm.types.guardrails import GuardrailEventHooks + + instances = [] + + def factory(**kwargs): + inst = Mock(close=AsyncMock()) + inst._kwargs = kwargs + instances.append(inst) + return inst + + install_sdk_stub(client_factory=factory) + + g = WonderFenceGuardrail( + guardrail_name="t", + api_key="default", + event_hook=[GuardrailEventHooks.pre_call], + ) + c1 = await g._get_client("key-A") + c1_again = await g._get_client("key-A") + c2 = await g._get_client("key-B") + assert c1 is c1_again + assert c1 is not c2 + assert len(instances) == 2 + + +@pytest.mark.asyncio +async def test_get_client_lru_evicts_oldest(install_sdk_stub): + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.alice_wonderfence import ( + WonderFenceGuardrail, + ) + from litellm.types.guardrails import GuardrailEventHooks + + def factory(**kwargs): + return Mock(close=AsyncMock(), _api_key=kwargs["api_key"]) + + install_sdk_stub(client_factory=factory) + + g = WonderFenceGuardrail( + guardrail_name="t", + api_key="default", + max_cached_clients=2, + event_hook=[GuardrailEventHooks.pre_call], + ) + a = await g._get_client("A") + b = await g._get_client("B") + # Touching A makes B the LRU candidate. + await g._get_client("A") + c = await g._get_client("C") # should evict B + + assert "A" in g._client_cache + assert "C" in g._client_cache + assert "B" not in g._client_cache + # Evicted client must NOT be closed — in-flight requests may still hold a + # reference. GC handles cleanup. + b.close.assert_not_awaited() + assert a is g._client_cache["A"] + assert c is g._client_cache["C"] + + +@pytest.mark.asyncio +async def test_get_client_forwards_config_to_v2_client(install_sdk_stub): + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.alice_wonderfence import ( + WonderFenceGuardrail, + ) + from litellm.types.guardrails import GuardrailEventHooks + + captured = [] + + def factory(**kwargs): + captured.append(kwargs) + return Mock(close=AsyncMock()) + + install_sdk_stub(client_factory=factory) + + g = WonderFenceGuardrail( + guardrail_name="t", + api_key="default", + api_base="https://wf.example.com", + api_timeout=15.4, + platform="aws", + connection_pool_limit=42, + event_hook=[GuardrailEventHooks.pre_call], + ) + await g._get_client("resolved-key") + + assert captured[0]["api_key"] == "resolved-key" + assert captured[0]["base_url"] == "https://wf.example.com" + assert captured[0]["api_timeout"] == 15 # rounded to int + assert captured[0]["connection_pool_limit"] == 42 + # ``platform`` must NOT be forwarded to the V2 client: its constructor is + # (api_key, base_url, *, api_timeout, connection_pool_limit) and has no such + # parameter, so forwarding it raised TypeError on every scan. platform is a + # per-request analysis attribute that belongs on AnalysisContext instead + # (see test_build_analysis_context_sets_platform_on_context). + assert "platform" not in captured[0] + + +# ----------------------------- initialization ----------------------------- + + +def test_initialization_falls_back_to_env(monkeypatch, make_guardrail): + monkeypatch.setenv("ALICE_API_KEY", "env-key") + guardrail, _ = make_guardrail(api_key=None) + assert guardrail.api_key == "env-key" + + +def test_initialization_no_default_api_key_does_not_raise(monkeypatch, make_guardrail): + """V2 model resolves api_key per-request — init must NOT require it.""" + monkeypatch.delenv("ALICE_API_KEY", raising=False) + guardrail, _ = make_guardrail(api_key=None) + assert guardrail.api_key is None + + +def test_allow_request_metadata_override_defaults_false(make_guardrail): + """New flag must default to False so request-body metadata cannot + bypass admin-pinned credentials out of the box.""" + guardrail, _ = make_guardrail() + assert guardrail.allow_request_metadata_override is False + + +def test_initialize_guardrail_forwards_all_params(install_sdk_stub): + """The package-level initializer must forward every typed config field.""" + install_sdk_stub() + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="alice_wonderfence", + mode="pre_call", + api_key="cfg-key", + api_base="https://wf.example.com", + api_timeout=12.0, + platform="aws", + fail_open=True, + block_message="custom block", + debug=True, + max_cached_clients=5, + connection_pool_limit=20, + allow_request_metadata_override=True, + default_on=True, + ) + guardrail = {"guardrail_name": "wf-init-test"} + + g = initialize_guardrail(params, guardrail) # type: ignore[arg-type] + + assert g.api_key == "cfg-key" + assert g.api_base == "https://wf.example.com" + assert g.api_timeout == 12.0 + assert g.platform == "aws" + assert g.fail_open is True + assert g.block_message == "custom block" + assert g._client_cache_maxsize == 5 + assert g._connection_pool_limit == 20 + assert g.allow_request_metadata_override is True + + +def test_initialize_guardrail_missing_name_raises(install_sdk_stub): + """Initializer rejects guardrails without a guardrail_name.""" + install_sdk_stub() + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + params = LitellmParams(guardrail="alice_wonderfence", mode="pre_call") + with pytest.raises(ValueError, match="requires a guardrail_name"): + initialize_guardrail(params, {}) # type: ignore[arg-type] + + +def test_init_raises_when_sdk_not_installed(monkeypatch): + """Constructor surfaces a clean ImportError when wonderfence_sdk missing.""" + monkeypatch.setitem(sys.modules, "wonderfence_sdk", None) + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.alice_wonderfence import ( + WonderFenceGuardrail, + ) + + with pytest.raises(ImportError, match="wonderfence-sdk"): + WonderFenceGuardrail(guardrail_name="t") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_credentials.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_credentials.py new file mode 100644 index 00000000000..9bb60c1c16d --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_credentials.py @@ -0,0 +1,328 @@ +"""Tests for credential resolution (api_key, app_id) helpers. + +These helpers are pure functions (no SDK dependency), so tests call them +directly with explicit args instead of constructing a guardrail instance. +""" + +import pytest + +import json + +from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.credentials import ( + get_metadata, + recover_resolved, + resolve_api_key, + resolve_app_id, + stash_resolved, +) +from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.exceptions import ( + WonderFenceMissingSecrets, +) + + +def _data(**overrides): + """Build a request_data dict with default admin-pinned app_id.""" + metadata = overrides.pop("metadata", None) + if metadata is None: + metadata = {"user_api_key_metadata": {"alice_wonderfence_app_id": "test-app"}} + base = {"model": "gpt-4", "metadata": metadata} + base.update(overrides) + return base + + +# ----------------------------- app_id resolution ----------------------------- + + +def test_resolve_app_id_from_request_metadata_requires_override_flag(): + data = _data(metadata={"alice_wonderfence_app_id": "from-req"}) + assert resolve_app_id(data, allow_request_metadata_override=True) == "from-req" + + +def test_resolve_app_id_request_metadata_ignored_when_override_disabled(): + """Request metadata is caller-controlled and must not satisfy app_id when + the override flag is off — otherwise a user could bypass admin-pinned + credentials by sending their own app_id in the request body.""" + data = _data(metadata={"alice_wonderfence_app_id": "from-req"}) + with pytest.raises(WonderFenceMissingSecrets, match="alice_wonderfence_app_id"): + resolve_app_id(data, allow_request_metadata_override=False) + + +def test_resolve_app_id_from_key_metadata(): + data = _data( + metadata={ + "user_api_key_metadata": {"alice_wonderfence_app_id": "from-key"}, + } + ) + assert resolve_app_id(data, allow_request_metadata_override=False) == "from-key" + + +def test_resolve_app_id_from_team_metadata(): + data = _data( + metadata={ + "user_api_key_team_metadata": {"alice_wonderfence_app_id": "from-team"}, + } + ) + assert resolve_app_id(data, allow_request_metadata_override=False) == "from-team" + + +def test_resolve_app_id_key_beats_request_even_when_override_enabled(): + """With the override flag on, request metadata is still only a last-resort + source — admin-pinned key metadata wins.""" + data = _data( + metadata={ + "alice_wonderfence_app_id": "from-req", + "user_api_key_metadata": {"alice_wonderfence_app_id": "from-key"}, + "user_api_key_team_metadata": {"alice_wonderfence_app_id": "from-team"}, + } + ) + assert resolve_app_id(data, allow_request_metadata_override=True) == "from-key" + + +def test_resolve_app_id_team_beats_request_when_override_enabled(): + """Team metadata beats request metadata even with the override flag on.""" + data = _data( + metadata={ + "alice_wonderfence_app_id": "from-req", + "user_api_key_team_metadata": {"alice_wonderfence_app_id": "from-team"}, + } + ) + assert resolve_app_id(data, allow_request_metadata_override=True) == "from-team" + + +def test_resolve_app_id_priority_key_over_team(): + data = _data( + metadata={ + "user_api_key_metadata": {"alice_wonderfence_app_id": "from-key"}, + "user_api_key_team_metadata": {"alice_wonderfence_app_id": "from-team"}, + } + ) + assert resolve_app_id(data, allow_request_metadata_override=False) == "from-key" + + +def test_resolve_app_id_missing_raises(): + data = _data(metadata={}) + with pytest.raises(WonderFenceMissingSecrets, match="alice_wonderfence_app_id"): + resolve_app_id(data, allow_request_metadata_override=False) + + +# ----------------------------- api_key resolution ----------------------------- + + +def test_resolve_api_key_from_request_metadata_requires_override_flag(): + data = _data(metadata={"alice_wonderfence_api_key": "from-req"}) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=True) == "from-req" + + +def test_resolve_api_key_request_metadata_ignored_when_override_disabled(): + """With override off, a caller-supplied api_key must not be honored; + falls back to the configured default instead.""" + data = _data(metadata={"alice_wonderfence_api_key": "from-req"}) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=False) == "default" + + +def test_resolve_api_key_key_beats_request_even_when_override_enabled(): + """Admin-pinned key metadata wins over request metadata even with the + override flag enabled.""" + data = _data( + metadata={ + "alice_wonderfence_api_key": "from-req", + "user_api_key_metadata": {"alice_wonderfence_api_key": "from-key"}, + } + ) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=True) == "from-key" + + +def test_resolve_api_key_from_key_metadata(): + data = _data( + metadata={ + "user_api_key_metadata": {"alice_wonderfence_api_key": "from-key"}, + } + ) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=False) == "from-key" + + +def test_resolve_api_key_from_team_metadata(): + data = _data( + metadata={ + "user_api_key_team_metadata": {"alice_wonderfence_api_key": "from-team"}, + } + ) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=False) == "from-team" + + +def test_resolve_api_key_falls_back_to_default(): + data = _data(metadata={}) + assert resolve_api_key(data, default_api_key="default-key", allow_request_metadata_override=False) == "default-key" + + +def test_resolve_api_key_missing_everywhere_raises(): + data = _data(metadata={}) + with pytest.raises(WonderFenceMissingSecrets): + resolve_api_key(data, default_api_key=None, allow_request_metadata_override=False) + + +# ----------------------------- metadata fallback ----------------------------- + + +def test_resolve_reads_litellm_metadata_when_metadata_absent(): + """``get_metadata`` falls back to ``litellm_metadata`` when ``metadata`` + is missing. Use admin-controlled key metadata so it resolves without + needing the request-override flag.""" + data = { + "model": "gpt-4", + "litellm_metadata": {"user_api_key_metadata": {"alice_wonderfence_app_id": "from-litellm-md"}}, + } + assert resolve_app_id(data, allow_request_metadata_override=False) == "from-litellm-md" + + +def test_get_metadata_merges_with_litellm_metadata_winning(): + """When both buckets are present, merge them with proxy-injected + ``litellm_metadata`` winning on key collision; caller-only keys survive.""" + data = { + "metadata": { + "alice_wonderfence_app_id": "caller-only", + "shared_key": "from-caller", + }, + "litellm_metadata": { + "shared_key": "from-litellm", + "user_api_key_metadata": {"alice_wonderfence_app_id": "admin-pinned"}, + }, + } + assert get_metadata(data) == { + "alice_wonderfence_app_id": "caller-only", + "shared_key": "from-litellm", + "user_api_key_metadata": {"alice_wonderfence_app_id": "admin-pinned"}, + } + + +def test_get_metadata_returns_empty_when_both_absent(): + assert get_metadata({}) == {} + + +def test_get_metadata_ignores_non_dict_caller_metadata(): + """A caller can send ``metadata`` as a non-object value. It must be coerced + away rather than returned verbatim, so the proxy-injected ``litellm_metadata`` + (carrying the admin pins) is preserved.""" + data = { + "metadata": "not-a-dict", + "litellm_metadata": {"user_api_key_metadata": {"alice_wonderfence_app_id": "admin-pinned"}}, + } + assert get_metadata(data) == {"user_api_key_metadata": {"alice_wonderfence_app_id": "admin-pinned"}} + + +def test_non_dict_caller_metadata_does_not_bypass_resolution(): + """Regression: a non-dict ``metadata`` once propagated to ``resolve_*`` and + raised on ``.get()``, which ``fail_open=True`` would swallow into a skipped + scan. Resolution must instead succeed from the admin pin in + ``litellm_metadata``.""" + data = { + "model": "gpt-4", + "metadata": ["unexpected", "list"], + "litellm_metadata": { + "user_api_key_metadata": { + "alice_wonderfence_app_id": "admin-pinned", + "alice_wonderfence_api_key": "admin-key", + } + }, + } + assert resolve_app_id(data, allow_request_metadata_override=True) == "admin-pinned" + assert resolve_api_key(data, default_api_key=None, allow_request_metadata_override=True) == "admin-key" + + +def test_responses_route_admin_pin_beats_caller_metadata(): + """Mirror the /v1/responses shape: caller `metadata` carries a + request-override app_id while the admin pin lives in + `litellm_metadata.user_api_key_metadata`. The admin pin must win even with + the override flag enabled — the caller bucket must not shadow it.""" + data = { + "model": "gpt-4", + "metadata": {"alice_wonderfence_app_id": "caller-override"}, + "litellm_metadata": {"user_api_key_metadata": {"alice_wonderfence_app_id": "admin-pinned"}}, + } + assert resolve_app_id(data, allow_request_metadata_override=True) == "admin-pinned" + + +def test_responses_route_admin_pin_beats_caller_metadata_api_key(): + """api_key variant of the /v1/responses regression: admin-pinned key + metadata wins over a caller-supplied request-override api_key.""" + data = { + "model": "gpt-4", + "metadata": {"alice_wonderfence_api_key": "caller-override"}, + "litellm_metadata": {"user_api_key_metadata": {"alice_wonderfence_api_key": "admin-pinned"}}, + } + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=True) == "admin-pinned" + + +# --------------- stash storage: secret must not leak to logged payload --------------- + + +def test_stash_round_trips_on_real_logging_obj(make_logging_obj): + """Guard: the real LiteLLMLoggingObj must accept the per-guardrail stash + attribute. Fails if Logging becomes slotted/pydantic.""" + obj = make_logging_obj() + stash_resolved(obj, "guard-1", "wf-key-abc", "app-1") + assert recover_resolved(obj, "guard-1") == ("wf-key-abc", "app-1") + + +def test_recover_does_not_borrow_a_sibling_guardrails_stash(make_logging_obj): + """Each guardrail's stash is isolated: a name that never stashed recovers + None even when a sibling stashed on the same logging_obj. This is what makes + a stricter instance fail closed instead of inheriting a permissive sibling's + caller-supplied credentials.""" + obj = make_logging_obj() + stash_resolved(obj, "writer", "writer-key", "writer-app") + assert recover_resolved(obj, "reader") is None + assert recover_resolved(obj, "writer") == ("writer-key", "writer-app") + + +def test_stashed_api_key_not_present_in_model_call_details(make_logging_obj): + """Regression: model_call_details is forwarded verbatim as kwargs to logging + callbacks/exporters, so a resolved tenant api_key stashed there leaks. The + stash must live off model_call_details. Fails on the prior implementation + that stored it under model_call_details["alice_wonderfence_resolved"].""" + secret = "wf-super-secret-key-9f3a" + obj = make_logging_obj() + stash_resolved(obj, "guard-1", secret, "app-1") + + dumped = json.dumps(obj.model_call_details, default=str) + assert secret not in dumped + assert "alice_wonderfence_resolved" not in obj.model_call_details + # recovery still works from the private attribute + assert recover_resolved(obj, "guard-1") == (secret, "app-1") + + +def test_recover_returns_none_when_nothing_stashed(make_logging_obj): + obj = make_logging_obj() + assert recover_resolved(obj, "guard-1") is None + + +# --------------- malformed (non-string) credential overrides --------------- + + +def test_resolve_api_key_ignores_non_string_request_override(): + """A truthy non-string request override must not be returned (it would reach + the SDK and raise, which fail_open could swallow); fall back to default.""" + data = _data(metadata={"alice_wonderfence_api_key": ["not", "a", "string"]}) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=True) == "default" + + +def test_resolve_app_id_non_string_request_override_raises(): + data = _data(metadata={"alice_wonderfence_app_id": {"bad": 1}}) + with pytest.raises(WonderFenceMissingSecrets): + resolve_app_id(data, allow_request_metadata_override=True) + + +def test_resolve_api_key_ignores_blank_string_override(): + data = _data(metadata={"alice_wonderfence_api_key": " "}) + assert resolve_api_key(data, default_api_key="default", allow_request_metadata_override=True) == "default" + + +def test_resolve_app_id_non_string_key_metadata_falls_through(): + """A non-string admin value is also rejected rather than passed to the SDK.""" + data = _data( + metadata={ + "user_api_key_metadata": {"alice_wonderfence_app_id": 12345}, + "user_api_key_team_metadata": {"alice_wonderfence_app_id": "team-app"}, + } + ) + assert resolve_app_id(data, allow_request_metadata_override=False) == "team-app" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_post_call_bridge.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_post_call_bridge.py new file mode 100644 index 00000000000..ce1849e2ccf --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_post_call_bridge.py @@ -0,0 +1,227 @@ +"""Tests for the post_call logging_obj stash + sibling fallback bridge.""" + +from unittest.mock import Mock + +import pytest +from fastapi import HTTPException + + +@pytest.mark.asyncio +async def test_post_call_recovers_app_id_via_logging_obj_stash(make_guardrail, make_request_data, make_logging_obj): + """Reproduces the framework gap: request body metadata is dropped before + post_call. The logging_obj stash from the prior ``input_type="request"`` + call must be used to resolve app_id.""" + guardrail, client = make_guardrail(allow_request_metadata_override=True) + guardrail._client_cache["default-api-key"] = client + request_obj = Mock() + request_obj.action = "NO_ACTION" + request_obj.detections = [] + request_obj.correlation_id = None + client.evaluate_prompt.return_value = request_obj + response_obj = Mock() + response_obj.action = "NO_ACTION" + response_obj.detections = [] + response_obj.correlation_id = None + client.evaluate_response.return_value = response_obj + + logging_obj = make_logging_obj() + + # Step 1: simulate pre_call / during_call with full request body + # metadata — this is where the stash happens. + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data=make_request_data(metadata={"alice_wonderfence_app_id": "tenant-X"}), + input_type="request", + logging_obj=logging_obj, + ) + + # Step 2: simulate post_call as the framework actually invokes it — + # the request body's metadata is gone (only litellm_metadata.user_api_key_* + # would normally be present, neither populated here). Without the + # bridge this raises; with it, we recover from logging_obj. + out = await guardrail.apply_guardrail( + inputs={"texts": ["llm response"]}, + request_data={"model": "gpt-4", "metadata": {}}, + input_type="response", + logging_obj=logging_obj, + ) + assert out["texts"] == ["llm response"] + assert client.evaluate_response.call_args.kwargs["app_id"] == "tenant-X" + + +@pytest.mark.asyncio +async def test_post_call_prefers_request_data_over_stash(make_guardrail, make_request_data, make_logging_obj): + """If post_call's request_data still resolves (e.g. app_id from key/team + metadata), use it — don't fall back to the stash.""" + guardrail, client = make_guardrail(allow_request_metadata_override=True) + guardrail._client_cache["default-api-key"] = client + request_obj = Mock() + request_obj.action = "NO_ACTION" + request_obj.detections = [] + request_obj.correlation_id = None + client.evaluate_prompt.return_value = request_obj + response_obj = Mock() + response_obj.action = "NO_ACTION" + response_obj.detections = [] + response_obj.correlation_id = None + client.evaluate_response.return_value = response_obj + + logging_obj = make_logging_obj() + + # Stash a different app_id during the request phase. + await guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={"alice_wonderfence_app_id": "stashed-app"}), + input_type="request", + logging_obj=logging_obj, + ) + + # Post_call request_data resolves via key metadata to a DIFFERENT app_id. + # The resolver path must win over the stash. + await guardrail.apply_guardrail( + inputs={"texts": ["resp"]}, + request_data={ + "model": "gpt-4", + "metadata": {"user_api_key_metadata": {"alice_wonderfence_app_id": "key-app"}}, + }, + input_type="response", + logging_obj=logging_obj, + ) + assert client.evaluate_response.call_args.kwargs["app_id"] == "key-app" + + +@pytest.mark.asyncio +async def test_post_call_without_prior_stash_raises(make_guardrail, make_logging_obj): + """If neither request_data nor logging_obj has the app_id (e.g. mode is + post_call only and app_id was supplied only in the request body), the + error path must still fire — not silently allow.""" + guardrail, client = make_guardrail() + guardrail._client_cache["default-api-key"] = client + + logging_obj = make_logging_obj() # empty model_call_details + + with pytest.raises(HTTPException) as exc: + await guardrail.apply_guardrail( + inputs={"texts": ["resp"]}, + request_data={"model": "gpt-4", "metadata": {}}, + input_type="response", + logging_obj=logging_obj, + ) + assert exc.value.status_code == 500 + assert "alice_wonderfence_app_id" in exc.value.detail["exception"] + + +@pytest.mark.asyncio +async def test_post_call_does_not_borrow_sibling_stash(make_guardrail, make_request_data, make_logging_obj): + """A stricter instance must NOT inherit a sibling's stashed credentials. + + Exploit being closed: a permissive writer (allow_request_metadata_override + =True) stashes caller-supplied request-body credentials; a stricter reader + (allow_request_metadata_override=False) that has no own stash must fail + closed in post_call rather than scan with the writer's caller-controlled + creds.""" + g_writer, c_writer = make_guardrail( + guardrail_name="writer", + allow_request_metadata_override=True, + ) + g_writer._client_cache["default-api-key"] = c_writer + g_reader, c_reader = make_guardrail( + guardrail_name="reader", + allow_request_metadata_override=False, + ) + g_reader._client_cache["default-api-key"] = c_reader + for c in (c_writer, c_reader): + result = Mock() + result.action = "NO_ACTION" + result.detections = [] + result.correlation_id = None + c.evaluate_prompt.return_value = result + c.evaluate_response.return_value = result + + logging_obj = make_logging_obj() + + # Writer stashes caller-supplied request-body app_id (override allowed). + await g_writer.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={"alice_wonderfence_app_id": "caller-supplied-app"}), + input_type="request", + logging_obj=logging_obj, + ) + + # Reader's post_call: no own stash, request_data resolves nothing, and it + # must NOT borrow the writer's stash -> fail closed. + with pytest.raises(HTTPException) as exc: + await g_reader.apply_guardrail( + inputs={"texts": ["resp"]}, + request_data={"model": "gpt-4", "metadata": {}}, + input_type="response", + logging_obj=logging_obj, + ) + assert exc.value.status_code == 500 + assert "alice_wonderfence_app_id" in exc.value.detail["exception"] + c_reader.evaluate_response.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stash_keyed_per_guardrail_name(make_guardrail, make_request_data, make_logging_obj): + """Two alice_wonderfence instances on the same logging_obj must not + overwrite each other's stash — they're keyed by guardrail_name.""" + g1, c1 = make_guardrail( + guardrail_name="alice-a", + allow_request_metadata_override=True, + ) + g1._client_cache["default-api-key"] = c1 + g2, c2 = make_guardrail( + guardrail_name="alice-b", + allow_request_metadata_override=True, + ) + g2._client_cache["default-api-key"] = c2 + for c in (c1, c2): + result = Mock() + result.action = "NO_ACTION" + result.detections = [] + result.correlation_id = None + c.evaluate_prompt.return_value = result + c.evaluate_response.return_value = result + + logging_obj = make_logging_obj() + + # Both instances stash under the SAME logging_obj using DIFFERENT + # request app_ids. + await g1.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={"alice_wonderfence_app_id": "app-a"}), + input_type="request", + logging_obj=logging_obj, + ) + await g2.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=make_request_data(metadata={"alice_wonderfence_app_id": "app-b"}), + input_type="request", + logging_obj=logging_obj, + ) + + # Each must recover its own value on post_call. + await g1.apply_guardrail( + inputs={"texts": ["resp"]}, + request_data={"model": "gpt-4", "metadata": {}}, + input_type="response", + logging_obj=logging_obj, + ) + await g2.apply_guardrail( + inputs={"texts": ["resp"]}, + request_data={"model": "gpt-4", "metadata": {}}, + input_type="response", + logging_obj=logging_obj, + ) + assert c1.evaluate_response.call_args.kwargs["app_id"] == "app-a" + assert c2.evaluate_response.call_args.kwargs["app_id"] == "app-b" + + +def test_recover_resolved_with_no_logging_obj_returns_none(): + """``recover_resolved`` must short-circuit on None logging_obj.""" + from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.credentials import ( + recover_resolved, + ) + + assert recover_resolved(None, "any-name") is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_processing.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_processing.py new file mode 100644 index 00000000000..204dfa384be --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/alice_wonderfence/test_processing.py @@ -0,0 +1,248 @@ +"""Tests for processing.py pure transforms: reconstruction, extractors, cap, verdict apply.""" + +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.chunked_evaluation import ( + SegmentVerdict, +) +from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.exceptions import ( + WonderFenceBlockedError, + WonderFenceScanBudgetExceeded, +) +from litellm.proxy.guardrails.guardrail_hooks.alice_wonderfence.processing import ( + JOINER, + RECONSTRUCT_MAX_CHARS, + apply_response_verdicts, + build_analysis_context, + check_scan_budget, + function_definition_segments, + reconstruct, + tool_definition_segments, +) + + +# --------------- build_analysis_context: platform belongs on the context --------------- + + +def test_build_analysis_context_sets_platform_on_context(): + """platform is a per-request analysis attribute and must be set on the + AnalysisContext, NOT forwarded to the SDK client constructor (whose + signature has no platform param). Pairs with + test_get_client_forwards_config_to_v2_client asserting platform is not + passed to the client.""" + captured: dict = {} + + def context_class(**kwargs): + captured.update(kwargs) + return object() + + build_analysis_context({"model": "gpt-4"}, "aws", context_class) + assert captured["platform"] == "aws" + + +def _block(detections=None, correlation_ids=None): + return SegmentVerdict("BLOCK", None, detections or [], correlation_ids or []) + + +# --------------- reconstruct (per-chunk masked alignment) --------------- +# +# ``masked_chunks`` is the list of (owned_original, owned_masked) regions that +# concatenate to the joined document; for a single-chunk (<= prompt limit) +# document that is just ``[(join, masked_join)]``. + + +def _one_chunk(parts, masked): + original = JOINER.join(parts) + return [(original, masked)] + + +def test_reconstruct_no_change_round_trips(): + parts = ["alpha", "beta", "gamma"] + assert reconstruct(parts, _one_chunk(parts, JOINER.join(parts))) == parts + + +def test_reconstruct_masks_a_middle_part(): + parts = ["alpha", "sensitive", "gamma"] + masked = JOINER.join(["alpha", "[REDACTED]", "gamma"]) + assert reconstruct(parts, _one_chunk(parts, masked)) == ["alpha", "[REDACTED]", "gamma"] + + +def test_reconstruct_mask_at_part_start(): + parts = ["alpha", "beta", "gamma"] + masked = JOINER.join(["[X]lpha", "beta", "gamma"]) + assert reconstruct(parts, _one_chunk(parts, masked)) == ["[X]lpha", "beta", "gamma"] + + +def test_reconstruct_handles_a_part_that_itself_contains_newline(): + """A message part can itself contain the joiner char; alignment is + structural, not a naive split on '\\n', so this still reconstructs.""" + parts = ["line1\nline1b", "second"] + masked = JOINER.join(["line1\n[REDACTED]", "second"]) + assert reconstruct(parts, _one_chunk(parts, masked)) == ["line1\n[REDACTED]", "second"] + + +def test_reconstruct_fails_closed_when_mask_spans_a_joiner(): + """If the mask swallows a joiner (parts merged), reconstruction must fail + closed (None) rather than misassign redacted text to the wrong message.""" + parts = ["alpha", "beta", "gamma"] + merged = "alphaXXXbeta\ngamma" # joiner between alpha|beta is gone + assert reconstruct(parts, _one_chunk(parts, merged)) is None + + +def test_reconstruct_maps_across_multiple_chunks(): + """The document is aligned per owned-region chunk; a part living in a later + chunk is recovered through that chunk's own alignment, not a global diff.""" + parts = ["aaaa", "bbbb"] + # Two owned regions that concatenate to "aaaa\nbbbb"; the second is masked. + masked_chunks = [("aaaa\n", "aaaa\n"), ("bbbb", "XXXX")] + assert reconstruct(parts, masked_chunks) == ["aaaa", "XXXX"] + + +def test_reconstruct_fails_closed_when_chunks_do_not_match_join(): + """Invariant guard: the owned regions must concatenate back to the join.""" + parts = ["alpha", "beta"] + assert reconstruct(parts, [("alpha\nDIFFERENT", "alpha\nDIFFERENT")]) is None + + +def test_reconstruct_none_chunks_fails_closed(): + assert reconstruct(["a", "b"], None) is None + + +def test_reconstruct_empty_parts_is_empty_list(): + assert reconstruct([], None) == [] + + +def test_reconstruct_fails_closed_when_document_exceeds_bound(): + """Reconstruction is bounded as a coarse backstop on total alignment work; + an over-bound document fails closed (None) instead of aligning.""" + big = "a" * (RECONSTRUCT_MAX_CHARS + 1) + assert reconstruct([big], [(big, big)]) is None + + +# --------------- check_scan_budget (total-work cap) --------------- + + +def test_check_scan_budget_passes_within_limits(): + check_scan_budget(["a", "b", "c"], max_scan_chars=100, max_scan_segments=100) + + +def test_check_scan_budget_rejects_too_many_segments(): + with pytest.raises(WonderFenceScanBudgetExceeded) as exc: + check_scan_budget(["x"] * 11, max_scan_chars=10_000, max_scan_segments=10) + assert exc.value.detail["limit"] == "max_scan_segments" + assert exc.value.detail["max_scan_segments"] == 10 + + +def test_check_scan_budget_rejects_too_many_chars(): + with pytest.raises(WonderFenceScanBudgetExceeded) as exc: + check_scan_budget(["x" * 50, "y" * 60], max_scan_chars=100, max_scan_segments=100) + assert exc.value.detail["limit"] == "max_scan_chars" + assert exc.value.detail["chars"] == 110 + + +def test_check_scan_budget_none_limits_disable_the_cap(): + check_scan_budget(["x" * 10_000] * 100, max_scan_chars=None, max_scan_segments=None) + + +# --------------- tool/function definition extractors (detection-only, list of texts) --------------- + + +def test_tool_definition_segments_extracts_description_and_param_descriptions(): + inputs = { + "tools": [ + { + "type": "function", + "function": { + "name": "weather", + "description": "TOP_DESC", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "PARAM_DESC"}}, + }, + }, + } + ] + } + assert set(tool_definition_segments(inputs)) == {"TOP_DESC", "PARAM_DESC"} + + +def test_tool_definition_segments_ignores_non_dict_tools_and_blank_descriptions(): + inputs = { + "tools": [ + "not-a-dict", + {"type": "function", "function": {"name": "f", "description": " "}}, + {"type": "function"}, + ] + } + assert tool_definition_segments(inputs) == [] + + +def test_function_definition_segments_extracts_descriptions(): + request_data = { + "functions": [ + { + "name": "weather", + "description": "TOP_DESC", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "PARAM_DESC"}}, + }, + }, + "not-a-dict", + {"name": "f", "description": " "}, + ] + } + assert set(function_definition_segments(request_data)) == {"TOP_DESC", "PARAM_DESC"} + + +def test_function_definition_segments_empty_when_absent(): + assert function_definition_segments({"model": "gpt-4"}) == [] + + +# --------------- apply_response_verdicts --------------- + + +def test_response_block_verdict_raises_with_aggregated_detections(): + inputs = {"texts": ["bad", "ok"]} + d = {"policy_name": "p"} + verdicts = [_block(detections=[d], correlation_ids=["c1"]), SegmentVerdict("", None, [], [])] + with pytest.raises(WonderFenceBlockedError) as exc: + apply_response_verdicts(inputs, verdicts, [], [], "gn", "blocked!") + assert exc.value.detail["error"] == "blocked!" + assert exc.value.detail["action"] == "BLOCK" + assert exc.value.detail["detections"] == [d] + assert exc.value.detail["wonderfence_correlation_id"] == "c1" + + +def test_response_mask_writes_to_the_mapped_text_index_only(): + inputs = {"texts": ["keep", "MASK_ME", "keep2"]} + verdicts = [ + SegmentVerdict("", None, [], []), + SegmentVerdict("MASK", "[R]", [], []), + SegmentVerdict("", None, [], []), + ] + out = apply_response_verdicts(inputs, verdicts, [], [], "gn", "blocked!") + assert out["texts"] == ["keep", "[R]", "keep2"] + + +def test_response_mask_writes_tool_call_arguments_in_place(): + inputs = { + "texts": ["ok"], + "tool_calls": [{"function": {"arguments": '{"x": "secret"}'}}], + } + out = apply_response_verdicts( + inputs, + [SegmentVerdict("", None, [], [])], + [0], + [SegmentVerdict("MASK", '{"x": "[R]"}', [], [])], + "gn", + "blocked!", + ) + assert out["tool_calls"][0]["function"]["arguments"] == '{"x": "[R]"}' + + +def test_response_detect_and_no_action_leave_texts_unchanged(): + inputs = {"texts": ["a", "b"]} + verdicts = [SegmentVerdict("DETECT", None, [], []), SegmentVerdict("", None, [], [])] + out = apply_response_verdicts(inputs, verdicts, [], [], "gn", "blocked!") + assert out["texts"] == ["a", "b"]