From c3ad8ac34e81e1861a84cdbf74e0a180e17cb865 Mon Sep 17 00:00:00 2001 From: Itay Ovadia Date: Fri, 14 Aug 2026 16:29:49 +0300 Subject: [PATCH] feat(guardrails): add payload, applicability, dispatch and logging controls to generic_guardrail_api Guardrail providers that are observe-only, that ignore parts of the payload, or that only care about some traffic currently have no way to say so: LiteLLM sends the whole payload for every call type, blocks the request until the guardrail answers, and appends a logging entry on every single call. On long agent sessions that means the entire transcript is re-sent every turn, synchronously, on the user's critical path, and one spend-log row accumulates thousands of guardrail entries. Seven additive options, all defaulting to today's behavior: - fire_and_forget dispatches the call as a background task and lets the request proceed, which makes the guardrail observe-only (BLOCKED and GUARDRAIL_INTERVENED can no longer be honored, and it says so at init). fire_and_forget_max_inflight bounds the queue so a slow endpoint cannot pile up tasks; in-flight tasks are strong-referenced so the event loop cannot collect them mid-call. Streaming is forced to end-of-stream observation, so a stream dispatches one call instead of one per sampled chunk. - send_images, exclude_payload_fields, max_messages and max_text_chars declare what the endpoint receives. send_images=false also replaces inline image data in structured_messages with a marker, so the base64 payload really does stay in LiteLLM. max_messages windows the flat texts list as well as structured_messages, since the translation handlers build texts from the whole conversation and windowing one without the other would leave payload size proportional to session length. input_type and litellm_call_id cannot be excluded. - strip_patterns removes configured regex matches from text before the payload is sent, touching only string text fields, never JSON structure, tool schemas, ids or metadata. - skip_if_system_prompt_matches and skip_if_first_role_in skip a call entirely, anchored to the system/developer message so pasted user content cannot trigger them. Both read the request body, which the caller controls, so a caller that knows the configured value can exempt itself; that is stated in the field descriptions and warned about at init. skip_if_key_alias_in and skip_if_team_id_in are the alternative that holds against the caller, matching on what authentication established. The message-based decision is replayed on the paired response through the shared logging object, with a bounded call-id cache for paths that have none; the identity filters need no correlation, since both hooks see the same metadata. - run_only_on_call_types and skip_call_types skip whole call types (embeddings, image, audio, rerank, moderation) before anything is sent. An unresolvable call type still runs, so a misconfiguration cannot silently blind the guardrail. - guardrail_information_scope picks how often a StandardLoggingGuardrailInformation entry is recorded: per_call keeps today's behavior, per_session records only the first call of a session, off records nothing on success. Session dedup needs a resolvable session id and falls back to per_call without one, and the session set is bounded by size and TTL. Suppression rides the log_guardrail_information contract via a new named entry point, suppress_guardrail_information_record, and is applied only after a normal return, so blocks and guardrail failures still record under every scope. Payload shaping is lossy, so each shaped payload carries the loss with it and the guardrail refuses to write back text the endpoint never saw in full: a stripped, truncated or windowed block keeps the caller's original content, and images or tools that were not sent cannot be replaced. Shaping runs on the dumped JSON because pydantic validates structured_messages content lazily into a ValidatorIterator that is unsafe to walk or re-iterate. The message-mapping helpers live in the shared guardrails content utils, next to the existing walkers, so the Chat Completions and Responses content shapes are handled in one place. --- litellm/integrations/custom_guardrail.py | 14 + litellm/proxy/guardrails/_content_utils.py | 147 +- .../generic_guardrail_api/__init__.py | 16 + .../background_dispatch.py | 85 ++ .../generic_guardrail_api.py | 322 ++++- .../generic_guardrail_api/payload_policy.py | 224 ++++ .../generic_guardrail_api/record_scope.py | 64 + .../generic_guardrail_api/request_filters.py | 187 +++ .../guardrail_hooks/generic_guardrail_api.py | 165 +++ .../test_generic_guardrail_api.py | 1194 +++++++++++++++++ .../proxy/guardrails/test_content_utils.py | 84 ++ 11 files changed, 2463 insertions(+), 39 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/background_dispatch.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/request_filters.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index f721e01e2c8..7dc9b37b5ac 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -65,6 +65,20 @@ _guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.Cont ) +def suppress_guardrail_information_record() -> None: + """Skip the automatic ``StandardLoggingGuardrailInformation`` entry for this call. + + ``log_guardrail_information`` clears the flag before invoking the guardrail and + skips its own append when the flag is set afterwards, so this is the supported + way for a guardrail to opt one invocation out of recording (for instance a + guardrail configured to record once per session rather than per call). + + Call it only on a success path: the decorator's exception branch reads the same + flag, so setting it before a raise also drops the error entry. + """ + _guardrail_self_recorded.set(True) + + def _strict_guardrail_modes_enabled() -> bool: """Whether guardrail-mode validation raises (default) or logs a warning. diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6ed6f0013df..635034cecb7 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -8,8 +8,10 @@ skip the other shapes — these helpers normalise that so every hook sees every text fragment. """ -from collections.abc import Callable, Iterator -from typing import Any, Final +from collections.abc import Callable, Collection, Iterator, Mapping, Sequence +from typing import Any, Final, TypeAlias + +from pydantic import JsonValue # Call types whose body carries free-form chat / prompt text that # text-content guardrails (banned keywords, content moderation, secret @@ -35,6 +37,12 @@ def is_text_content_call_type(call_type: str) -> bool: TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "output_text"}) +# A text rewrite applied to one fragment at a time. +TextTransform: TypeAlias = Callable[[str], str] + +# Content part types whose payload is image data (Chat Completions and Responses API). +IMAGE_PART_TYPES: Final[frozenset[str]] = frozenset({"image_url", "input_image"}) + # Responses-API item types whose ``output`` field carries user/tool text # that guardrails should inspect. ``function_call_output`` is the # built-in shape; ``custom_tool_call_output`` is the custom-tool @@ -104,6 +112,111 @@ def iter_message_text(data: dict[str, Any]) -> Iterator[str]: yield from _iter_text_parts_in_content(message.get("content")) +def iter_role_text(messages: Sequence[Mapping[str, object]], roles: Collection[str]) -> Iterator[str]: + """Yield every text fragment carried by messages whose role is in ``roles``. + + Lets a hook anchor a decision to a specific role (e.g. the system / + developer instructions) instead of the whole transcript, which would match + on any user-pasted content. + """ + for message in messages: + # The annotation is the contract; the values arrive from provider + # translation handlers, so a stray non-message entry must not raise. + if not isinstance(message, Mapping): # pyright: ignore[reportUnnecessaryIsInstance] # untyped upstream data + continue + if message.get("role") in roles: + yield from _iter_text_parts_in_content(message.get("content")) + + +def map_content_text(content: JsonValue, transform: TextTransform) -> JsonValue: + """Return ``content`` with every text fragment replaced by ``transform(fragment)``. + + Pure counterpart of ``walk_user_text``: string content, bare strings in a + content list, and ``{"type": "text", "text": ...}`` parts are transformed; + every other part (images, audio, tool payloads) is passed through by + reference so the content structure cannot change shape. + """ + if isinstance(content, str): + return transform(content) if content else content + if not isinstance(content, list): + return content + return [_mapped_text_part(part, transform) for part in content] # mutable-ok: JSON content is a list + + +def _mapped_text_part(part: JsonValue, transform: TextTransform) -> JsonValue: + if isinstance(part, str): + # A bare string in a content list is itself a text fragment. + return transform(part) if part else part + if not isinstance(part, dict) or part.get("type") not in TEXT_PART_TYPES: + return part + text: Final = part.get("text") + if not isinstance(text, str) or not text: + return part + return {**part, "text": transform(text)} # mutable-ok: JSON parts are dicts + + +def map_content_image_urls(content: JsonValue, transform: TextTransform) -> JsonValue: + """Return ``content`` with every image payload replaced by ``transform(url)``. + + Covers both image part shapes: Chat Completions ``{"type": "image_url", + "image_url": {"url": ...}}`` (and its bare-string variant) and the + Responses-API ``{"type": "input_image", "image_url": ...}``. Everything else + is passed through by reference. + """ + if not isinstance(content, list): + return content + return [_mapped_image_part(part, transform) for part in content] # mutable-ok: JSON content is a list + + +def _mapped_image_part(part: JsonValue, transform: TextTransform) -> JsonValue: + if not isinstance(part, dict) or part.get("type") not in IMAGE_PART_TYPES: + return part + image_url: Final = part.get("image_url") + if isinstance(image_url, str): + return {**part, "image_url": transform(image_url)} # mutable-ok: JSON parts are dicts + if not isinstance(image_url, dict): + return part + url: Final = image_url.get("url") + if not isinstance(url, str): + return part + nested: Final = {**image_url, "url": transform(url)} # mutable-ok: JSON parts are dicts + return {**part, "image_url": nested} # mutable-ok: JSON parts are dicts + + +def map_messages_image_urls(messages: JsonValue, transform: TextTransform) -> JsonValue: + """Return a new message list with every image payload transformed.""" + return _mapped_messages(messages, map_content_image_urls, transform) + + +def map_messages_text(messages: JsonValue, transform: TextTransform) -> JsonValue: + """Return a new message list with every ``content`` text fragment transformed. + + Roles, ids, tool calls, tool schemas and any other key are copied through + untouched, so only text can differ from the input. + """ + return _mapped_messages(messages, map_content_text, transform) + + +def _mapped_messages( + messages: JsonValue, + map_content: Callable[[JsonValue, TextTransform], JsonValue], + transform: TextTransform, +) -> JsonValue: + if not isinstance(messages, list): + return messages + return [_mapped_message(message, map_content, transform) for message in messages] # mutable-ok: JSON is a list + + +def _mapped_message( + message: JsonValue, + map_content: Callable[[JsonValue, TextTransform], JsonValue], + transform: TextTransform, +) -> JsonValue: + if not isinstance(message, dict) or "content" not in message: + return message + return {**message, "content": map_content(message["content"], transform)} # mutable-ok: JSON messages are dicts + + def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: """Rewrite every text fragment in place via ``visit``. @@ -113,31 +226,13 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: """ visited = 0 - def _rewrite_content(content: Any) -> Any: + def _counting_visit(text: str) -> str: nonlocal visited - if isinstance(content, str): - if content: - visited += 1 - return visit(content) - return content - if isinstance(content, list): - new_parts: Final[list[Any]] = [] - for part in content: - if isinstance(part, str) and part: - visited += 1 - new_parts.append(visit(part)) - elif ( - isinstance(part, dict) - and part.get("type") in TEXT_PART_TYPES - and isinstance(part.get("text"), str) - and part["text"] - ): - visited += 1 - new_parts.append({**part, "text": visit(part["text"])}) - else: - new_parts.append(part) - return new_parts - return content + visited += 1 + return visit(text) + + def _rewrite_content(content: Any) -> Any: + return map_content_text(content, _counting_visit) messages: Final = data.get("messages") if isinstance(messages, list): diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index d1576b68813..8cb9e95eeec 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -39,6 +39,22 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"), streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"), streaming_transform_mode=_get_config_value(litellm_params, optional_params, "streaming_transform_mode"), + fire_and_forget=_get_config_value(litellm_params, optional_params, "fire_and_forget"), + fire_and_forget_max_inflight=_get_config_value(litellm_params, optional_params, "fire_and_forget_max_inflight"), + send_images=_get_config_value(litellm_params, optional_params, "send_images"), + exclude_payload_fields=_get_config_value(litellm_params, optional_params, "exclude_payload_fields"), + max_messages=_get_config_value(litellm_params, optional_params, "max_messages"), + max_text_chars=_get_config_value(litellm_params, optional_params, "max_text_chars"), + strip_patterns=_get_config_value(litellm_params, optional_params, "strip_patterns"), + skip_if_system_prompt_matches=_get_config_value( + litellm_params, optional_params, "skip_if_system_prompt_matches" + ), + skip_if_first_role_in=_get_config_value(litellm_params, optional_params, "skip_if_first_role_in"), + skip_if_key_alias_in=_get_config_value(litellm_params, optional_params, "skip_if_key_alias_in"), + skip_if_team_id_in=_get_config_value(litellm_params, optional_params, "skip_if_team_id_in"), + run_only_on_call_types=_get_config_value(litellm_params, optional_params, "run_only_on_call_types"), + skip_call_types=_get_config_value(litellm_params, optional_params, "skip_call_types"), + guardrail_information_scope=_get_config_value(litellm_params, optional_params, "guardrail_information_scope"), ) litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/background_dispatch.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/background_dispatch.py new file mode 100644 index 00000000000..a4bc41bde68 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/background_dispatch.py @@ -0,0 +1,85 @@ +"""Off-critical-path dispatch for the Generic Guardrail API (``fire_and_forget``). + +The event loop only holds a weak reference to a task created with +``asyncio.create_task``, so a task with no other referrer can be garbage +collected mid-flight and silently cancelled. Every dispatched task is therefore +kept in a strong-ref set until it completes. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from typing import Final + +from litellm._logging import verbose_proxy_logger + +DEFAULT_FIRE_AND_FORGET_MAX_INFLIGHT: Final = 100 + +# Log one warning per this many drops, so a saturated endpoint cannot turn into a +# log flood of its own. +_DROP_LOG_INTERVAL: Final = 100 + + +class BackgroundDispatcher: + """Runs guardrail calls detached from the request, with a bounded queue. + + Async dispatch decouples the request rate from the endpoint's throughput, so + a slow endpoint would otherwise pile up tasks without limit. Once + ``max_inflight`` tasks are outstanding, further calls are dropped and + counted rather than queued. + """ + + def __init__(self, *, guardrail_name: str | None, max_inflight: int) -> None: + if max_inflight < 1: + raise ValueError(f"fire_and_forget_max_inflight must be >= 1 (got {max_inflight})") + self._guardrail_name: Final = guardrail_name + self._max_inflight: Final = max_inflight + # Strong refs to in-flight tasks; entries are discarded on completion. + self._pending: set[asyncio.Task[None]] = set() # mutable-ok: tasks are added/removed as they complete + self._dropped: int = 0 + + @property + def pending_count(self) -> int: + return len(self._pending) + + @property + def dropped_count(self) -> int: + return self._dropped + + def dispatch(self, run: Callable[[], Awaitable[None]], *, context: str) -> bool: + """Start ``run`` detached. Returns False when dropped for backpressure.""" + if len(self._pending) >= self._max_inflight: + self._dropped += 1 + if self._dropped % _DROP_LOG_INTERVAL == 1: + verbose_proxy_logger.warning( + "Generic Guardrail API (%s, fire_and_forget): dropped %d call(s) so far; " + "%d already in flight (fire_and_forget_max_inflight=%d). %s", + self._guardrail_name, + self._dropped, + len(self._pending), + self._max_inflight, + context, + ) + return False + + task: Final = asyncio.create_task(self._guarded(run, context=context)) + self._pending.add(task) + task.add_done_callback(self._pending.discard) + return True + + async def _guarded(self, run: Callable[[], Awaitable[None]], *, context: str) -> None: + """Never raise: a detached task's exception has no frame to surface in.""" + try: + await run() + except Exception as e: # noqa: BLE001 # a detached task has no frame to raise into + verbose_proxy_logger.warning( + "Generic Guardrail API (%s, fire_and_forget) call failed. %s: %s", + self._guardrail_name, + context, + e, + ) + + async def wait_for_pending(self) -> None: + """Await every in-flight task. Used by tests and best-effort teardown.""" + pending: Final = tuple(self._pending) + if pending: + await asyncio.gather(*pending, return_exceptions=True) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16768a4b08f..797c252bdaf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -7,22 +7,27 @@ import fnmatch import os +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional import httpx +from pydantic import JsonValue from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.exceptions import GuardrailRaisedException, Timeout from litellm.integrations.custom_guardrail import ( CustomGuardrail, + get_session_id_from_request_data, log_guardrail_information, + suppress_guardrail_information_record, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, @@ -31,6 +36,32 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ) from litellm.types.utils import GenericGuardrailAPIInputs +from .background_dispatch import ( + DEFAULT_FIRE_AND_FORGET_MAX_INFLIGHT, + BackgroundDispatcher, +) +from .payload_policy import ( + PayloadLoss, + PayloadPolicy, + compile_patterns, + merge_guardrailed_texts, + resolve_exclude_fields, + shape_payload, +) +from .record_scope import ( + DEFAULT_GUARDRAIL_INFORMATION_SCOPE, + GuardrailInformationScope, + RecordScope, +) +from .request_filters import ( + SkipDecisionStore, + SkipPolicy, + call_type_allowed, + identity_matches_skip, + request_matches_skip, + validate_call_types, +) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -148,6 +179,42 @@ def _extract_inbound_headers( return None +def _resolve_call_type( + request_data: Mapping[str, object], + logging_obj: Optional["LiteLLMLoggingObj"], +) -> str | None: + """Resolve the call type of the current request. + + The proxy passes the route type straight through to ``function_setup``, which + stamps it on the logging object before the pre-call hook runs and keeps it + for the post-call hook, so both sides of a call resolve the same value. + """ + from_logging: Final = getattr(logging_obj, "call_type", None) if logging_obj else None + if isinstance(from_logging, str) and from_logging: + return from_logging + from_request: Final = request_data.get("call_type") + return from_request if isinstance(from_request, str) and from_request else None + + +def _passthrough_inputs(inputs: GenericGuardrailAPIInputs) -> GenericGuardrailAPIInputs: + """Return the inputs untouched (same value identities), as action=NONE.""" + return GenericGuardrailAPIInputs(**inputs) + + +def _has_request_side_hook(event_hook: str | Sequence[str] | Mode | None) -> bool: + """Whether the configured mode(s) include a hook that sees the request. + + Tag-based ``Mode`` config is resolved per request, so it is treated as having + one rather than emitting a warning that may not apply. + """ + request_side: Final = frozenset({GuardrailEventHooks.pre_call.value, GuardrailEventHooks.during_call.value}) + if event_hook is None or isinstance(event_hook, Mode): + return True + if isinstance(event_hook, str): + return event_hook in request_side + return any(hook in request_side for hook in event_hook) + + class GenericGuardrailAPI(CustomGuardrail): """ Generic Guardrail API integration for LiteLLM. @@ -182,9 +249,27 @@ class GenericGuardrailAPI(CustomGuardrail): streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None, + fire_and_forget: bool | None = None, + fire_and_forget_max_inflight: int | None = None, + send_images: bool | None = None, + exclude_payload_fields: Sequence[str] | None = None, + max_messages: int | None = None, + max_text_chars: int | None = None, + strip_patterns: Sequence[str] | None = None, + skip_if_system_prompt_matches: Sequence[str] | None = None, + skip_if_first_role_in: Sequence[str] | None = None, + skip_if_key_alias_in: Sequence[str] | None = None, + skip_if_team_id_in: Sequence[str] | None = None, + run_only_on_call_types: Sequence[str] | None = None, + skip_call_types: Sequence[str] | None = None, + guardrail_information_scope: GuardrailInformationScope | None = None, + async_handler: AsyncHTTPHandler | None = None, + dispatcher: BackgroundDispatcher | None = None, **kwargs, ): - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.headers = headers or {} self.extra_headers = extra_headers or [] @@ -229,6 +314,93 @@ class GenericGuardrailAPI(CustomGuardrail): "block_only" if streaming_transform_mode is None else streaming_transform_mode ) + configured_name: Final = kwargs.get("guardrail_name") + + self.fire_and_forget: bool = False if fire_and_forget is None else fire_and_forget + self._dispatcher: Final = dispatcher or BackgroundDispatcher( + guardrail_name=configured_name, + max_inflight=( + DEFAULT_FIRE_AND_FORGET_MAX_INFLIGHT + if fire_and_forget_max_inflight is None + else fire_and_forget_max_inflight + ), + ) + + self._payload_policy: Final = PayloadPolicy( + send_images=True if send_images is None else send_images, + exclude_fields=resolve_exclude_fields(exclude_payload_fields, guardrail_name=configured_name), + max_messages=max_messages, + max_text_chars=max_text_chars, + strip_patterns=compile_patterns(strip_patterns, option_name="strip_patterns"), + ) + + self._skip_policy: Final = SkipPolicy( + system_prompt_patterns=compile_patterns( + skip_if_system_prompt_matches, option_name="skip_if_system_prompt_matches" + ), + first_role_in=frozenset(skip_if_first_role_in or ()), + key_aliases=frozenset(skip_if_key_alias_in or ()), + team_ids=frozenset(skip_if_team_id_in or ()), + run_only_on_call_types=( + validate_call_types( + run_only_on_call_types, + option_name="run_only_on_call_types", + guardrail_name=configured_name, + ) + if run_only_on_call_types + else None + ), + skip_call_types=validate_call_types( + skip_call_types, option_name="skip_call_types", guardrail_name=configured_name + ), + ) + self._skip_store: Final = SkipDecisionStore(guardrail_name=configured_name) + + self._record_scope: Final = RecordScope( + DEFAULT_GUARDRAIL_INFORMATION_SCOPE if guardrail_information_scope is None else guardrail_information_scope + ) + + if self.fire_and_forget: + # Nothing awaits the response, so a block cannot be honored and the + # stream can only be observed. Say so at boot rather than surprising + # the operator with a guardrail that never blocks. + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): fire_and_forget=True makes this guardrail observe-only. " + "action=BLOCKED and action=GUARDRAIL_INTERVENED are ignored, and " + "fail_on_error=%s / unreachable_fallback=%s cannot block the request. " + "Streaming is forced to end-of-stream observation.", + configured_name, + self.fail_on_error, + self.unreachable_fallback, + ) + self.streaming_end_of_stream_only = True + + if self._skip_policy.filters_requests: + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): skip_if_system_prompt_matches / skip_if_first_role_in match on the " + "request body, which the caller controls, so a caller that knows the configured value can exempt " + "itself from this guardrail. Use skip_if_key_alias_in / skip_if_team_id_in when the exemption must " + "hold against the caller.", + configured_name, + ) + + if self._skip_policy.filters_requests and not _has_request_side_hook(kwargs.get("event_hook")): + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): skip_if_system_prompt_matches / skip_if_first_role_in need a " + "request-side hook (pre_call or during_call) to decide anything. mode=%s only sees " + "responses, so nothing will be skipped.", + configured_name, + kwargs.get("event_hook"), + ) + + if self._skip_policy.run_only_on_call_types is not None and self._skip_policy.skip_call_types: + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): both run_only_on_call_types and skip_call_types are set. " + "The allowlist wins; skip_call_types=%s is ignored.", + configured_name, + sorted(self._skip_policy.skip_call_types), + ) + # Set supported event hooks kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -303,9 +475,7 @@ class GenericGuardrailAPI(CustomGuardrail): exc_info=error, ) # Keep flow going - treat as action=NONE (no modifications) - return_inputs: Final[GenericGuardrailAPIInputs] = {} - return_inputs.update(inputs) - return return_inputs + return _passthrough_inputs(inputs) def _build_request_headers(self) -> dict: """Build HTTP headers for the guardrail API request.""" @@ -321,16 +491,24 @@ class GenericGuardrailAPI(CustomGuardrail): images: Any, tools: Any, guardrail_response: GenericGuardrailAPIResponse, + loss: PayloadLoss, ) -> GenericGuardrailAPIInputs: - # Action is NONE or no modifications needed + # Action is NONE or no modifications needed. A component the guardrail + # never received in full (payload shaping) keeps the caller's original + # value: it cannot rewrite what it could not see. return_inputs: Final = GenericGuardrailAPIInputs(texts=texts) if guardrail_response.texts: - return_inputs["texts"] = guardrail_response.texts - if guardrail_response.images: + return_inputs["texts"] = merge_guardrailed_texts( + original=texts, + returned=guardrail_response.texts, + loss=loss, + guardrail_name=getattr(self, "guardrail_name", None), + ) + if guardrail_response.images and not loss.images_omitted: return_inputs["images"] = guardrail_response.images elif images: return_inputs["images"] = images - if guardrail_response.tools: + if guardrail_response.tools and not loss.tools_omitted: return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools @@ -359,6 +537,54 @@ class GenericGuardrailAPI(CustomGuardrail): verbose_proxy_logger.error("Generic Guardrail API: failed to make request: %s", str(error)) raise Exception(f"Generic Guardrail API failed: {error}") + def _should_skip_out_of_scope_request( + self, + *, + input_type: Literal["request", "response"], + structured_messages: Sequence[Mapping[str, object]] | None, + request_data: Mapping[str, object], + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> bool: + """Whether this call is out of scope per skip_if_* (Feature 4). + + Only the request carries the system prompt, so the response side replays + the decision the request side recorded instead of re-deciding. + """ + if not self._skip_policy.filters_requests: + return False + + call_id: Final = (getattr(logging_obj, "litellm_call_id", None) if logging_obj else None) or request_data.get( + "litellm_call_id" + ) + + if input_type == "response": + return self._skip_store.consume(logging_obj=logging_obj, call_id=call_id) + + if not request_matches_skip(self._skip_policy, structured_messages): + return False + + self._skip_store.record(logging_obj=logging_obj, call_id=call_id) + return True + + def _dispatch_background_post( + self, + *, + payload: Mapping[str, JsonValue], + headers: Mapping[str, str], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> None: + context: Final = ( + f"input_type={input_type} " + f"litellm_call_id={getattr(logging_obj, 'litellm_call_id', None) if logging_obj else None}" + ) + + async def _post() -> None: + response: Final = await self.async_handler.post(url=self.api_base, json=payload, headers=headers) + response.raise_for_status() + + self._dispatcher.dispatch(_post, context=context) + @log_guardrail_information async def apply_guardrail( self, @@ -366,6 +592,29 @@ class GenericGuardrailAPI(CustomGuardrail): request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Run the guardrail, then decide whether this call records a log entry. + + The suppression flag is set only once the call has returned normally, so + a block or a guardrail failure still records under every scope: the + decorator's exception branch reads the same flag. + """ + result: Final = await self._apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=logging_obj, + ) + if self._record_scope.should_suppress(get_session_id_from_request_data(request_data or {})): + suppress_guardrail_information_record() + return result + + async def _apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: """ Apply the Generic Guardrail API to the given inputs. @@ -401,6 +650,40 @@ class GenericGuardrailAPI(CustomGuardrail): if request_data is None: request_data = {} + # Extract user API key metadata (also the basis of the identity filters below) + user_metadata: Final = self._extract_user_api_key_metadata(request_data) + + call_type: Final = _resolve_call_type(request_data=request_data, logging_obj=logging_obj) + if not call_type_allowed(self._skip_policy, call_type): + verbose_proxy_logger.debug( + "Generic Guardrail API: skipping call_type=%s (input_type=%s) per call-type filter", + call_type, + input_type, + ) + return _passthrough_inputs(inputs) + + if identity_matches_skip(self._skip_policy, user_metadata): + verbose_proxy_logger.debug( + "Generic Guardrail API: skipping out-of-scope caller (input_type=%s, key_alias=%s, team_id=%s)", + input_type, + user_metadata.get("user_api_key_alias"), + user_metadata.get("user_api_key_team_id"), + ) + return _passthrough_inputs(inputs) + + if self._should_skip_out_of_scope_request( + input_type=input_type, + structured_messages=structured_messages, + request_data=request_data, + logging_obj=logging_obj, + ): + verbose_proxy_logger.debug( + "Generic Guardrail API: skipping out-of-scope request (input_type=%s, litellm_call_id=%s)", + input_type, + getattr(logging_obj, "litellm_call_id", None) if logging_obj else None, + ) + return _passthrough_inputs(inputs) + request_body: Final = request_data.get("body") or {} # Merge additional provider specific params from config and dynamic params @@ -411,8 +694,6 @@ class GenericGuardrailAPI(CustomGuardrail): if dynamic_params: additional_params.update(dynamic_params) - # Extract user API key metadata - user_metadata: Final = self._extract_user_api_key_metadata(request_data) extra_allowlist = {h.lower() for h in self.extra_headers if isinstance(h, str)} if self.extra_headers else None inbound_headers: Final = _extract_inbound_headers( request_data=request_data, @@ -440,11 +721,25 @@ class GenericGuardrailAPI(CustomGuardrail): headers: Final = self._build_request_headers() + # Shape the payload (send_images / exclude_payload_fields / max_messages / + # max_text_chars / strip_patterns) once, so the awaited and the + # fire_and_forget paths send exactly the same bytes. + # mode="json" ensures all iterables are converted to lists. + payload, payload_loss = shape_payload(guardrail_request, self._payload_policy) + + if self.fire_and_forget: + self._dispatch_background_post( + payload=payload, + headers=headers, + input_type=input_type, + logging_obj=logging_obj, + ) + return _passthrough_inputs(inputs) + # Make the API request - # Use mode="json" to ensure all iterables are converted to lists response: Final = await self.async_handler.post( url=self.api_base, - json=guardrail_request.model_dump(mode="json"), + json=payload, headers=headers, ) @@ -471,6 +766,7 @@ class GenericGuardrailAPI(CustomGuardrail): images=images, tools=tools, guardrail_response=guardrail_response, + loss=payload_loss, ) except GuardrailRaisedException: diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py new file mode 100644 index 00000000000..7b3e1c014dc --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py @@ -0,0 +1,224 @@ +"""Payload shaping for the Generic Guardrail API. + +Lets the config declare what the guardrail endpoint actually receives: +``send_images``, ``exclude_payload_fields``, ``max_messages``, +``max_text_chars`` and ``strip_patterns``. + +Shaping is lossy by design, so every shaped payload carries a ``PayloadLoss`` +describing what the endpoint could not see. The guardrail refuses to write back +mutations for those components: a provider that never received a text block (or +received a stripped / truncated one) must not be able to replace the caller's +original content with it. +""" + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import reduce +from typing import Final + +from pydantic import JsonValue + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.guardrails._content_utils import ( + map_messages_image_urls, + map_messages_text, +) +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIRequest, +) + +# Routing / correlation keys the endpoint needs to interpret any payload at all. +PROTECTED_PAYLOAD_FIELDS: Final = frozenset({"input_type", "litellm_call_id"}) + +# Ceiling on substitutions per pattern per text, so a pathological pattern on a +# huge transcript cannot stall the request path. +MAX_STRIP_SUBSTITUTIONS: Final = 64 + +# Stands in for image data when send_images=False, mirroring the "[present]" +# convention used for headers whose value is not forwarded. +IMAGE_OMITTED_PLACEHOLDER: Final = "[omitted]" + +_IMAGES_FIELD: Final = frozenset({"images"}) + + +@dataclass(frozen=True, slots=True) +class PayloadPolicy: + """Resolved payload-shaping configuration.""" + + send_images: bool = True + exclude_fields: frozenset[str] = frozenset() + max_messages: int | None = None + max_text_chars: int | None = None + strip_patterns: tuple[re.Pattern[str], ...] = () + + @property + def shapes_text(self) -> bool: + return self.max_text_chars is not None or bool(self.strip_patterns) + + +@dataclass(frozen=True, slots=True) +class PayloadLoss: + """What the guardrail endpoint did not get to see.""" + + altered_text_indices: frozenset[int] = frozenset() + images_omitted: bool = False + tools_omitted: bool = False + + +def resolve_exclude_fields(raw: Sequence[str] | None, *, guardrail_name: str | None) -> frozenset[str]: + """Validate ``exclude_payload_fields`` against the request model. + + Unknown keys and routing-critical keys are dropped with a warning rather + than raising, so a stale config never takes the proxy down. + """ + if not raw: + return frozenset() + + known: Final = frozenset(GenericGuardrailAPIRequest.model_fields) + unknown: Final = tuple(field for field in raw if field not in known) + if unknown: + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): exclude_payload_fields contains unknown field(s) %s; ignoring them. " + "Known fields: %s", + guardrail_name, + unknown, + sorted(known), + ) + + protected: Final = tuple(field for field in raw if field in PROTECTED_PAYLOAD_FIELDS) + if protected: + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): exclude_payload_fields cannot drop routing-critical field(s) %s; " + "they will still be sent.", + guardrail_name, + protected, + ) + + return frozenset(field for field in raw if field in known and field not in PROTECTED_PAYLOAD_FIELDS) + + +def compile_patterns(raw: Sequence[str] | None, *, option_name: str) -> tuple[re.Pattern[str], ...]: + """Compile config-supplied regexes, raising on an invalid pattern. + + A bad regex is a config bug: failing at init surfaces it at proxy boot + instead of silently disabling the option on the request path. + """ + if not raw: + return () + try: + return tuple(re.compile(pattern) for pattern in raw) + except re.error as e: + raise ValueError(f"{option_name} contains an invalid regex: {e}") from e + + +def _strip_text(text: str, patterns: tuple[re.Pattern[str], ...]) -> str: + return reduce(lambda acc, pattern: pattern.sub("", acc, count=MAX_STRIP_SUBSTITUTIONS), patterns, text) + + +def _shape_text(text: str, policy: PayloadPolicy) -> str: + stripped: Final = _strip_text(text, policy.strip_patterns) + return stripped[: policy.max_text_chars] if policy.max_text_chars is not None else stripped + + +def _shape_texts(texts: Sequence[str], policy: PayloadPolicy) -> tuple[JsonValue, frozenset[int]]: + """Window and rewrite the flat text list, and report which indices changed. + + ``max_messages`` has to bound ``texts`` as well as ``structured_messages``: + the translation handlers populate the flat list from the whole conversation, + so windowing only the structured form would leave payload size proportional + to session length, which is the thing the option exists to stop. Windowing + shifts every position, so the caller treats the whole list as unusable for + write-back (see ``merge_guardrailed_texts``). + """ + windowed: Final = texts[-policy.max_messages :] if policy.max_messages is not None else texts + shaped: Final = [_shape_text(text, policy) for text in windowed] # mutable-ok: JSON texts is an array + if len(windowed) != len(texts): + return shaped, frozenset(range(len(texts))) + return shaped, frozenset(index for index, text in enumerate(texts) if text != shaped[index]) + + +def _shape_messages(messages: JsonValue, policy: PayloadPolicy) -> JsonValue: + if not isinstance(messages, list) or not messages: + return None + windowed: Final = messages[-policy.max_messages :] if policy.max_messages is not None else messages + texted: Final = ( + map_messages_text(windowed, lambda text: _shape_text(text, policy)) if policy.shapes_text else windowed + ) + # send_images=False has to cover the inline image parts too, otherwise the + # base64 payload still leaves LiteLLM inside structured_messages. The part is + # kept so the guardrail still knows an image was there. + if policy.send_images: + return texted + return map_messages_image_urls(texted, lambda _: IMAGE_OMITTED_PLACEHOLDER) + + +def shape_payload( + request: GenericGuardrailAPIRequest, + policy: PayloadPolicy, +) -> tuple[Mapping[str, JsonValue], PayloadLoss]: + """Return the JSON payload to POST plus the loss it represents. + + Shaping runs on the dumped payload rather than on the model: pydantic + validates ``structured_messages`` list content lazily into a + ``ValidatorIterator``, which is neither a list to walk nor safe to iterate + twice, so the plain JSON form is the only reliable input here. + """ + exclude: Final = policy.exclude_fields if policy.send_images else policy.exclude_fields | _IMAGES_FIELD + excluded: Final = set(exclude) or None # mutable-ok: model_dump(exclude=...) takes a set + dumped: Final = request.model_dump(mode="json", exclude=excluded) + + texts: Final = request.texts + texts_sent: Final = bool(texts) and "texts" not in exclude + shaped_texts, altered_indices = _shape_texts(texts or (), policy) if texts_sent else (None, frozenset[int]()) + shaped_messages: Final = _shape_messages(dumped.get("structured_messages"), policy) + + payload: Final = { # mutable-ok: the POST body is a JSON object + **dumped, + **({"texts": shaped_texts} if texts_sent else {}), # mutable-ok: JSON object + **({"structured_messages": shaped_messages} if shaped_messages is not None else {}), # mutable-ok: JSON + } + loss: Final = PayloadLoss( + # An excluded text block was never seen at all, so no index may be rewritten. + altered_text_indices=(altered_indices if texts_sent else frozenset(range(len(texts or ())))), + images_omitted="images" in exclude, + tools_omitted="tools" in exclude, + ) + return payload, loss + + +def merge_guardrailed_texts( + *, + original: list[str], # mutable-ok: the framework's write-back contract + returned: list[str], # mutable-ok: the framework's write-back contract + loss: PayloadLoss, + guardrail_name: str | None, +) -> list[str]: # mutable-ok: GenericGuardrailAPIInputs["texts"] is list[str] + """Accept guardrail-returned text only where the endpoint saw the original. + + Indices the payload altered (stripped / truncated) or never carried keep the + caller's text, so lossy shaping can never splice a mangled prompt back into + the request. Length mismatch breaks index alignment, so with any loss the + whole rewrite is refused. + """ + if not loss.altered_text_indices: + return returned + if len(returned) != len(original): + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): ignoring returned texts. The payload was shaped " + "(strip_patterns / max_text_chars / excluded texts) and the guardrail returned " + "%d text(s) for %d sent, so they cannot be aligned.", + guardrail_name, + len(returned), + len(original), + ) + return original + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): keeping original text for shaped index(es) %s; the guardrail " + "never saw the full content at those positions.", + guardrail_name, + sorted(loss.altered_text_indices), + ) + return [ # mutable-ok: write-back contract + original[index] if index in loss.altered_text_indices else text for index, text in enumerate(returned) + ] diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py new file mode 100644 index 00000000000..28a9430bb9a --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py @@ -0,0 +1,64 @@ +"""How often this guardrail records a spend-log / OTEL entry. + +``@log_guardrail_information`` appends a ``StandardLoggingGuardrailInformation`` +entry on every invocation, with no cap and no dedup, so a long-lived session +accumulates one entry per guardrail call in the same request metadata. This +module decides, per call, whether that entry should be recorded at all. + +Suppression rides the decorator's existing contract: it clears +``_guardrail_self_recorded`` before calling the guardrail and skips its own +append if the flag is set afterwards. Setting the flag only on a normal return +(never before a raise) keeps blocks and guardrail failures recorded under every +scope, since the decorator's exception branch honors the same flag. +""" + +from typing import Final, Literal + +from litellm.caching.in_memory_cache import InMemoryCache + +GuardrailInformationScope = Literal["per_call", "per_session", "off"] + +DEFAULT_GUARDRAIL_INFORMATION_SCOPE: Final[GuardrailInformationScope] = "per_call" + +# Bounded so a long-running proxy cannot accumulate session keys forever. +_SESSION_CACHE_MAX_ENTRIES: Final = 100_000 +_SESSION_CACHE_TTL_SECONDS: Final = 3600 + + +class RecordScope: + """Tracks which sessions have already recorded an entry for one guardrail.""" + + def __init__( + self, + scope: GuardrailInformationScope, + *, + cache: InMemoryCache | None = None, + ) -> None: + self._scope: Final = scope + self._recorded_sessions: Final = cache or InMemoryCache( + max_size_in_memory=_SESSION_CACHE_MAX_ENTRIES, + default_ttl=_SESSION_CACHE_TTL_SECONDS, + ) + + @property + def scope(self) -> GuardrailInformationScope: + return self._scope + + def should_suppress(self, session_id: str | None) -> bool: + """Whether this successful call should skip its logging entry. + + ``per_session`` claims the session on its first call, so the calls that + follow it suppress. A session id is required: without one there is + nothing to dedup against, so the call records as it would under + ``per_call`` rather than silently dropping every entry. + """ + if self._scope == "per_call": + return False + if self._scope == "off": + return True + if session_id is None: + return False + if self._recorded_sessions.get_cache(session_id) is True: + return True + self._recorded_sessions.set_cache(session_id, True, ttl=_SESSION_CACHE_TTL_SECONDS) + return False diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/request_filters.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/request_filters.py new file mode 100644 index 00000000000..2f6882b5b2e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/request_filters.py @@ -0,0 +1,187 @@ +"""Applicability filters for the Generic Guardrail API. + +Two independent ways to keep a request away from the guardrail endpoint +entirely: + +- by call type (``run_only_on_call_types`` / ``skip_call_types``): every hook + resolves its own call type, so request and response skip symmetrically with no + correlation needed. +- by request content (``skip_if_system_prompt_matches`` / + ``skip_if_first_role_in``): only the request side can decide, so the decision + is recorded and replayed on the paired response via ``SkipDecisionStore``. +""" + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.guardrails._content_utils import iter_role_text +from litellm.types.utils import CallTypes + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +# Roles carrying the operator's instructions. Matching these instead of the whole +# transcript avoids skipping a request because a user pasted a matching string. +INSTRUCTION_ROLES: Final = frozenset({"system", "developer"}) + +# Fallback correlation cache, used only when the request and response hooks do +# not share a logging object. Bounded and short-lived: a response that never +# fires (stream abort, upstream error) must not leak an entry. +_SKIP_CACHE_MAX_ENTRIES: Final = 1000 +_SKIP_CACHE_TTL_SECONDS: Final = 300 + + +@dataclass(frozen=True, slots=True) +class SkipPolicy: + """Resolved applicability configuration.""" + + system_prompt_patterns: tuple[re.Pattern[str], ...] = () + first_role_in: frozenset[str] = frozenset() + key_aliases: frozenset[str] = frozenset() + team_ids: frozenset[str] = frozenset() + run_only_on_call_types: frozenset[str] | None = None + skip_call_types: frozenset[str] = frozenset() + + @property + def filters_requests(self) -> bool: + """Whether any message-based filter is configured. + + These read the request body, so only the request hook can decide and the + decision has to be replayed on the response. + """ + return bool(self.system_prompt_patterns) or bool(self.first_role_in) + + @property + def filters_identities(self) -> bool: + """Whether any caller-identity filter is configured. + + These read what authentication established, which both hooks see, so each + side decides for itself and no correlation is needed. + """ + return bool(self.key_aliases) or bool(self.team_ids) + + +def validate_call_types( + raw: Sequence[str] | None, + *, + option_name: str, + guardrail_name: str | None, +) -> frozenset[str]: + """Validate configured call types against ``CallTypes``, warning on unknowns.""" + if not raw: + return frozenset() + known: Final = frozenset(call_type.value for call_type in CallTypes) + unknown: Final = tuple(value for value in raw if value not in known) + if unknown: + verbose_proxy_logger.warning( + "Generic Guardrail API (%s): %s contains unrecognized call type(s) %s. " + "Values must be CallTypes names (e.g. acompletion, aembedding, anthropic_messages).", + guardrail_name, + option_name, + unknown, + ) + return frozenset(raw) + + +def call_type_allowed(policy: SkipPolicy, call_type: str | None) -> bool: + """Whether the guardrail runs for this call type. + + An unresolved call type runs, so a filter misconfiguration can never + silently blind the guardrail. The allowlist takes precedence over the + denylist. + """ + if call_type is None: + return True + if policy.run_only_on_call_types is not None: + return call_type in policy.run_only_on_call_types + return call_type not in policy.skip_call_types + + +def identity_matches_skip(policy: SkipPolicy, metadata: Mapping[str, object]) -> bool: + """Whether the calling key or team is out of scope for the guardrail. + + Matched against the metadata the auth layer produced, not the request body, + so a caller cannot exempt itself by choosing what it sends. + """ + if not policy.filters_identities: + return False + key_alias: Final = metadata.get("user_api_key_alias") + if policy.key_aliases and isinstance(key_alias, str) and key_alias in policy.key_aliases: + return True + team_id: Final = metadata.get("user_api_key_team_id") + return bool(policy.team_ids and isinstance(team_id, str) and team_id in policy.team_ids) + + +def request_matches_skip( + policy: SkipPolicy, + structured_messages: Sequence[Mapping[str, object]] | None, +) -> bool: + """Whether this request is out of scope for the guardrail.""" + if not policy.filters_requests or not structured_messages: + return False + + first_role: Final = structured_messages[0].get("role") + if policy.first_role_in and isinstance(first_role, str) and first_role in policy.first_role_in: + return True + + if not policy.system_prompt_patterns: + return False + + return any( + pattern.search(text) + for text in iter_role_text(structured_messages, INSTRUCTION_ROLES) + for pattern in policy.system_prompt_patterns + ) + + +class SkipDecisionStore: + """Carries a request-side skip decision to the paired response hook. + + Primary channel is the shared logging object, which is exact and dies with + the call. The bounded cache keyed by ``litellm_call_id`` covers the paths + where no logging object reaches the hooks; entries are evicted on read. + """ + + def __init__(self, *, guardrail_name: str | None, cache: InMemoryCache | None = None) -> None: + self._detail_key: Final = f"_generic_guardrail_skip::{guardrail_name}" + self._cache: Final = cache or InMemoryCache( + max_size_in_memory=_SKIP_CACHE_MAX_ENTRIES, + default_ttl=_SKIP_CACHE_TTL_SECONDS, + ) + + def _cache_key(self, call_id: str) -> str: + return f"{self._detail_key}::{call_id}" + + @staticmethod + def _call_details( + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> dict[str, bool] | None: # mutable-ok: this is the logging object's live dict, written through on purpose + """The shared write-through channel, or None when this call has no logging object.""" + details: Final[object] = getattr(logging_obj, "model_call_details", None) if logging_obj else None + # model_call_details is a free-form dict; only our own bool marker is read + # or written here, so the narrower value type is accurate for this use. + return details if isinstance(details, dict) else None + + def record(self, *, logging_obj: Optional["LiteLLMLoggingObj"], call_id: str | None) -> None: + details: Final = self._call_details(logging_obj) + if details is not None: + details[self._detail_key] = True + return + if call_id: + self._cache.set_cache(self._cache_key(call_id), True, ttl=_SKIP_CACHE_TTL_SECONDS) + + def consume(self, *, logging_obj: Optional["LiteLLMLoggingObj"], call_id: str | None) -> bool: + details: Final = self._call_details(logging_obj) + if details is not None and details.get(self._detail_key) is True: + return True + if not call_id: + return False + key: Final = self._cache_key(call_id) + if self._cache.get_cache(key) is not True: + return False + self._cache.delete_cache(key) + return True diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4a868c48352..d043f20506f 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -102,6 +102,171 @@ class GenericGuardrailAPIOptionalParams(BaseModel): ), ) + fire_and_forget: bool | None = Field( + default=None, + description=( + "If True, the guardrail HTTP call is dispatched as a background task and the " + "request proceeds immediately without awaiting the response. Applies to every " + "mode (pre_call, post_call, during_call), adding ~0 latency. Because the " + "response is never awaited, the guardrail is observe-only: action=BLOCKED and " + "action=GUARDRAIL_INTERVENED are ignored. Also forces " + "streaming_end_of_stream_only=True so a stream dispatches one background call " + "instead of one per sampled chunk. Defaults to False in " + "GenericGuardrailAPI.__init__ when None." + ), + ) + + fire_and_forget_max_inflight: int | None = Field( + default=None, + ge=1, + description=( + "Maximum number of fire_and_forget calls in flight at once. Async dispatch " + "decouples the request rate from the guardrail endpoint's throughput, so a slow " + "endpoint would otherwise pile up tasks without bound. Calls beyond this limit " + "are dropped and counted, with a rate-limited warning. Ignored unless " + "fire_and_forget is True. Defaults to 100 in GenericGuardrailAPI.__init__ when None." + ), + ) + + send_images: bool | None = Field( + default=None, + description=( + "If False, base64 image data is omitted from the guardrail request even when the " + "LLM request contains images. Large payload saver for guardrails that only " + "inspect text. Images the guardrail never received cannot be rewritten by its " + "response. Defaults to True in GenericGuardrailAPI.__init__ when None." + ), + ) + + exclude_payload_fields: tuple[str, ...] | None = Field( + default=None, + description=( + "Top-level GenericGuardrailAPIRequest keys to omit from the payload (e.g. " + "['images','texts','request_headers']), for providers that do not consume them. " + "Unknown keys and the routing-critical keys input_type and litellm_call_id are " + "ignored with a warning at init. A component that is not sent cannot be " + "rewritten by the guardrail response." + ), + ) + + max_messages: int | None = Field( + default=None, + ge=1, + description=( + "If set, only the last N entries of structured_messages and the last N text " + "blocks are sent. Bounds payload size when the full conversation is re-sent " + "every turn. Note the system prompt and early context fall out of the window " + "once the session exceeds N. LOSSY: windowing shifts text positions, so the " + "guardrail can no longer rewrite text on this call (action=BLOCKED still " + "applies); leave it unset on a guardrail that masks or redacts." + ), + ) + + max_text_chars: int | None = Field( + default=None, + ge=1, + description=( + "If set, each individual text block is truncated to this many characters before " + "sending, for providers that only need a prefix. LOSSY: truncated text blocks " + "keep their original content on write-back, so the guardrail cannot rewrite them." + ), + ) + + strip_patterns: tuple[str, ...] | None = Field( + default=None, + description=( + "Regex patterns applied to text content (texts[] and the text of each " + "structured_messages entry) before the guardrail request is sent. Matches are " + "removed. Intended to drop volatile boilerplate the provider does not need. " + "Applied only to string text fields, never to JSON structure, tool schemas, ids " + "or metadata. LOSSY: stripped content cannot be inspected by the guardrail, and " + "stripped text blocks keep their original content on write-back. An invalid " + "regex raises at init." + ), + ) + + skip_if_system_prompt_matches: tuple[str, ...] | None = Field( + default=None, + description=( + "Regex patterns matched against the request's system message (role=system or " + "developer). On a match the guardrail is skipped for this call: no request is " + "sent, and the paired response is skipped too, so both telemetry and enforcement " + "are suppressed for matched requests. Matching the system message only (not " + "arbitrary user text) avoids false positives from pasted content. Requires a " + "request-side hook (pre_call or during_call) in the mode list. " + "TRUST BOUNDARY: the system message comes from the request body, so any caller " + "who knows the configured pattern can add a system message and exempt itself " + "from this guardrail. Treat this as traffic scoping, not as enforcement, and " + "prefer skip_if_key_alias_in / skip_if_team_id_in when the exemption has to " + "hold against the caller." + ), + ) + + skip_if_first_role_in: tuple[str, ...] | None = Field( + default=None, + description=( + "If the first message's role is in this list (e.g. ['developer']), skip the " + "guardrail for the call, with the same request/response semantics as, and the " + "same trust boundary as, skip_if_system_prompt_matches: the caller chooses the " + "roles it sends." + ), + ) + + skip_if_key_alias_in: tuple[str, ...] | None = Field( + default=None, + description=( + "If the calling virtual key's alias is in this list, skip the guardrail for the " + "call. Unlike the message-based filters this reads what authentication " + "established, so a caller cannot exempt itself by changing its request body. " + "Same request/response semantics: neither side is sent." + ), + ) + + skip_if_team_id_in: tuple[str, ...] | None = Field( + default=None, + description=( + "If the calling key's team id is in this list, skip the guardrail for the call. " + "Admin-controlled like skip_if_key_alias_in, and matched on the same " + "authenticated metadata." + ), + ) + + run_only_on_call_types: tuple[str, ...] | None = Field( + default=None, + description=( + "If set, the guardrail runs ONLY for these call types (e.g. " + "['completion','acompletion','anthropic_messages','responses','aresponses']). All " + "other call types, including embeddings, image generation, audio, rerank and " + "moderation, are skipped before any request is sent. Allowlist; takes precedence " + "over skip_call_types. Values are CallTypes names. An unresolvable call type runs." + ), + ) + + skip_call_types: tuple[str, ...] | None = Field( + default=None, + description=( + "If set (and run_only_on_call_types is not), the guardrail is skipped for these " + "call types (e.g. ['embedding','aembedding','image_generation','transcription'," + "'speech','rerank','moderation']). Denylist. Values are CallTypes names." + ), + ) + + guardrail_information_scope: Literal["per_call", "per_session", "off"] | None = Field( + default=None, + description=( + "How often this guardrail records a StandardLoggingGuardrailInformation entry " + "into request metadata (spend logs / OTEL). Every invocation records one today, " + "with no cap and no dedup, so a long agent session accumulates one entry per " + "guardrail call in the same row. 'per_call' (default) keeps that behavior. " + "'per_session' records only the first call of a session, and needs a resolvable " + "session id (litellm_session_id or metadata.session_id); without one it behaves " + "as 'per_call' rather than dropping every entry. 'off' records nothing on " + "success. Blocks and guardrail failures are recorded under every scope. Dedup " + "is per proxy process, so a session spread across workers records once per " + "worker. Defaults to 'per_call' in GenericGuardrailAPI.__init__ when None." + ), + ) + class GenericGuardrailAPIConfigModel( GuardrailConfigModel[GenericGuardrailAPIOptionalParams], diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 5be0d43c250..a417be5e69b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -5,6 +5,7 @@ This test file tests the Generic Guardrail API implementation, specifically focusing on metadata extraction and passing. """ +import json as json_module import os from unittest.mock import AsyncMock, MagicMock, patch @@ -19,6 +20,9 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPI, ) +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.background_dispatch import ( + BackgroundDispatcher, +) from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( _HEADER_PRESENT_PLACEHOLDER, ) @@ -1996,3 +2000,1193 @@ class TestFailOnError: request_data={}, input_type="response", ) + + +# --------------------------------------------------------------------------- +# Payload shaping, applicability filters and fire-and-forget dispatch +# --------------------------------------------------------------------------- + + +class _RecordingHandler: + """Stands in for AsyncHTTPHandler, capturing every posted payload. + + Injected through GenericGuardrailAPI(async_handler=...) so the guardrail + under test keeps its real code path and nothing has to be patched onto it. + """ + + def __init__(self, *, action="NONE", texts=None, images=None, tools=None, error=None): + self.calls: list[dict] = [] + self._action = action + self._texts = texts + self._images = images + self._tools = tools + self._error = error + + @property + def payloads(self) -> list[dict]: + return [call["json"] for call in self.calls] + + async def post(self, *, url, json, headers, **kwargs): + self.calls.append({"url": url, "json": json, "headers": headers}) + if self._error is not None: + raise self._error + body = {"action": self._action} + for key, value in (("texts", self._texts), ("images", self._images), ("tools", self._tools)): + if value is not None: + body[key] = value + response = MagicMock() + response.json.return_value = body + response.raise_for_status = MagicMock() + return response + + +class _StubLoggingObj: + """Minimal stand-in for the LiteLLM logging object shared by both hooks.""" + + def __init__(self, *, call_type=None, call_id="call-123", trace_id="trace-123"): + self.call_type = call_type + self.litellm_call_id = call_id + self.litellm_trace_id = trace_id + self.model_call_details: dict = {} + + +def _make_guardrail(handler, *, dispatcher=None, name="test-generic-guardrail", event_hook="pre_call", **options): + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name=name, + event_hook=event_hook, + default_on=True, + async_handler=handler, + dispatcher=dispatcher, + **options, + ) + + +class TestCallTypeFilter: + """run_only_on_call_types / skip_call_types (call-type applicability).""" + + @pytest.mark.asyncio + async def test_allowlist_runs_allowed_call_type(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, run_only_on_call_types=["acompletion"]) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_type="acompletion"), + ) + + assert len(handler.calls) == 1 + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("input_type", ["request", "response"]) + async def test_allowlist_skips_embeddings_on_both_hooks(self, input_type): + """An embedding call is skipped on request and response alike; each hook + resolves its own call type, so no correlation is needed.""" + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, run_only_on_call_types=["acompletion"]) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["embed me"]}, + request_data={}, + input_type=input_type, + logging_obj=_StubLoggingObj(call_type="aembedding"), + ) + + assert handler.calls == [] + assert result == {"texts": ["embed me"]} + + @pytest.mark.asyncio + async def test_unresolved_call_type_still_runs(self): + """A call type we cannot resolve must not silently blind the guardrail.""" + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, run_only_on_call_types=["acompletion"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_type=None), + ) + + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_denylist_skips_listed_call_type_only(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, skip_call_types=["aembedding", "aspeech"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["embed me"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_type="aembedding"), + ) + assert handler.calls == [] + + await guardrail.apply_guardrail( + inputs={"texts": ["chat"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_type="acompletion"), + ) + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_allowlist_takes_precedence_over_denylist(self): + handler = _RecordingHandler() + guardrail = _make_guardrail( + handler, + run_only_on_call_types=["aembedding"], + skip_call_types=["aembedding"], + ) + + await guardrail.apply_guardrail( + inputs={"texts": ["embed me"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_type="aembedding"), + ) + + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_call_type_from_request_data_when_no_logging_obj(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, run_only_on_call_types=["acompletion"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["embed me"]}, + request_data={"call_type": "aembedding"}, + input_type="request", + ) + + assert handler.calls == [] + + +class TestPayloadFieldControl: + """send_images / exclude_payload_fields / max_messages / max_text_chars.""" + + @pytest.mark.asyncio + async def test_send_images_false_omits_images_key(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, send_images=False) + + await guardrail.apply_guardrail( + inputs={"texts": ["describe"], "images": ["data:image/png;base64,AAAA"]}, + request_data={}, + input_type="request", + ) + + payload = handler.payloads[0] + assert "images" not in payload + assert payload["texts"] == ["describe"] + + @pytest.mark.asyncio + async def test_send_images_false_also_strips_inline_image_parts(self): + """Regression: dropping the images array is not enough, structured_messages + carries the same base64 payload inline.""" + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, send_images=False) + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,SECRETPIXELS"}}, + ], + } + ] + + await guardrail.apply_guardrail( + inputs={ + "texts": ["describe this"], + "images": ["data:image/png;base64,SECRETPIXELS"], + "structured_messages": messages, + }, + request_data={}, + input_type="request", + ) + + payload = handler.payloads[0] + assert "SECRETPIXELS" not in json_module.dumps(payload) + # The part is kept, so the guardrail still sees that an image was sent. + part = payload["structured_messages"][0]["content"][1] + assert part["type"] == "image_url" + assert part["image_url"]["url"] == "[omitted]" + assert payload["structured_messages"][0]["content"][0]["text"] == "describe this" + + @pytest.mark.asyncio + async def test_bare_string_image_url_part_is_covered(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, send_images=False) + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": "data:image/png;base64,SECRETPIXELS"}, + ], + } + ] + + await guardrail.apply_guardrail( + inputs={"texts": ["describe this"], "structured_messages": messages}, + request_data={}, + input_type="request", + ) + + assert "SECRETPIXELS" not in json_module.dumps(handler.payloads[0]) + + @pytest.mark.asyncio + async def test_images_sent_by_default(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler) + + await guardrail.apply_guardrail( + inputs={"texts": ["describe"], "images": ["data:image/png;base64,AAAA"]}, + request_data={}, + input_type="request", + ) + + assert handler.payloads[0]["images"] == ["data:image/png;base64,AAAA"] + + @pytest.mark.asyncio + async def test_omitted_images_cannot_be_rewritten_by_guardrail(self): + """The guardrail never received the image, so its replacement is refused.""" + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", images=["data:image/png;base64,EVIL"]) + guardrail = _make_guardrail(handler, send_images=False) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["describe"], "images": ["data:image/png;base64,AAAA"]}, + request_data={}, + input_type="request", + ) + + assert result["images"] == ["data:image/png;base64,AAAA"] + + @pytest.mark.asyncio + async def test_exclude_payload_fields_drops_only_requested_field(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, exclude_payload_fields=["request_headers", "litellm_version"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"proxy_server_request": {"headers": {"user-agent": "curl/8"}}}, + input_type="request", + logging_obj=_StubLoggingObj(), + ) + + payload = handler.payloads[0] + assert "request_headers" not in payload + assert "litellm_version" not in payload + assert payload["texts"] == ["hello"] + assert payload["input_type"] == "request" + + @pytest.mark.asyncio + async def test_protected_fields_cannot_be_excluded(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, exclude_payload_fields=["input_type", "litellm_call_id"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_id="call-abc"), + ) + + payload = handler.payloads[0] + assert payload["input_type"] == "request" + assert payload["litellm_call_id"] == "call-abc" + + @pytest.mark.asyncio + async def test_max_messages_sends_only_the_tail(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, max_messages=2) + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "two"}, + {"role": "user", "content": "three"}, + ] + + await guardrail.apply_guardrail( + inputs={"texts": ["three"], "structured_messages": messages}, + request_data={}, + input_type="request", + ) + + sent = handler.payloads[0]["structured_messages"] + assert [m["content"] for m in sent] == ["two", "three"] + + @pytest.mark.asyncio + async def test_max_text_chars_truncates_each_text_block(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, max_text_chars=5) + + await guardrail.apply_guardrail( + inputs={"texts": ["0123456789", "ab"]}, + request_data={}, + input_type="request", + ) + + assert handler.payloads[0]["texts"] == ["01234", "ab"] + + @pytest.mark.asyncio + async def test_truncated_text_is_not_written_back(self): + """Regression: a rewrite of truncated text must not replace the full prompt.""" + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED", "ab"]) + guardrail = _make_guardrail(handler, max_text_chars=5) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["0123456789", "ab"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["0123456789", "ab"] + + @pytest.mark.asyncio + async def test_untouched_text_index_still_accepts_rewrite(self): + """Only the shaped index is protected; other indices are still masked.""" + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED-LONG", "MASKED-SHORT"]) + guardrail = _make_guardrail(handler, max_text_chars=5) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["0123456789", "ab"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["0123456789", "MASKED-SHORT"] + + @pytest.mark.asyncio + async def test_rewrite_refused_when_lengths_do_not_align(self): + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED"]) + guardrail = _make_guardrail(handler, max_text_chars=5) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["0123456789", "ab"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["0123456789", "ab"] + + @pytest.mark.asyncio + async def test_unshaped_payload_still_applies_rewrites(self): + """Default config keeps today's masking behavior intact.""" + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED"]) + guardrail = _make_guardrail(handler) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["MASKED"] + + @pytest.mark.asyncio + async def test_excluded_texts_are_neither_sent_nor_rewritten(self): + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED"]) + guardrail = _make_guardrail(handler, exclude_payload_fields=["texts"]) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123"]}, + request_data={}, + input_type="request", + ) + + assert "texts" not in handler.payloads[0] + assert result["texts"] == ["my ssn is 123"] + + def test_unknown_exclude_field_does_not_break_init(self): + guardrail = _make_guardrail(_RecordingHandler(), exclude_payload_fields=["not_a_field"]) + assert guardrail._payload_policy.exclude_fields == frozenset() + + +class TestStripPatterns: + """strip_patterns: source-side removal of content the provider does not need.""" + + ENV_BLOCK = "CWD=/tmp\nDATE=2026-01-01" + + @pytest.mark.asyncio + async def test_matched_content_never_leaves_litellm(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, strip_patterns=[r"[\s\S]*?"]) + + await guardrail.apply_guardrail( + inputs={"texts": [f"{self.ENV_BLOCK}\nreal question"]}, + request_data={}, + input_type="request", + ) + + sent = handler.payloads[0]["texts"][0] + assert "CWD=/tmp" not in sent + assert "real question" in sent + + @pytest.mark.asyncio + async def test_structured_messages_keep_their_structure(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, strip_patterns=[r"SECRET-\d+"]) + messages = [ + {"role": "system", "content": "you are SECRET-42 helpful"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "leak SECRET-7 here"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "run", "arguments": '{"cmd": "SECRET-9"}'}, + } + ], + }, + ] + + await guardrail.apply_guardrail( + inputs={"texts": ["leak SECRET-7 here"], "structured_messages": messages}, + request_data={}, + input_type="request", + ) + + sent = handler.payloads[0]["structured_messages"] + assert [m["role"] for m in sent] == ["system", "user", "assistant"] + assert sent[0]["content"] == "you are helpful" + assert sent[1]["content"][0]["text"] == "leak here" + # Non-text parts, tool calls and ids are never rewritten. + assert sent[1]["content"][1]["image_url"] == {"url": "data:image/png;base64,AAAA"} + assert sent[2]["tool_calls"][0]["function"]["arguments"] == '{"cmd": "SECRET-9"}' + assert sent[2]["tool_calls"][0]["id"] == "call_1" + + @pytest.mark.asyncio + async def test_tool_schemas_are_never_stripped(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, strip_patterns=[r"SECRET-\d+"]) + + await guardrail.apply_guardrail( + inputs={ + "texts": ["hello"], + "tools": [ + { + "type": "function", + "function": {"name": "SECRET-1", "description": "SECRET-2"}, + } + ], + }, + request_data={}, + input_type="request", + ) + + sent_tool = handler.payloads[0]["tools"][0] + assert sent_tool["function"]["name"] == "SECRET-1" + assert sent_tool["function"]["description"] == "SECRET-2" + + @pytest.mark.asyncio + async def test_stripped_text_is_not_written_back(self): + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["rewritten"]) + guardrail = _make_guardrail(handler, strip_patterns=[r"SECRET-\d+"]) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["keep SECRET-1 me"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["keep SECRET-1 me"] + + @pytest.mark.asyncio + async def test_non_matching_text_is_untouched(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, strip_patterns=[r"SECRET-\d+"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["nothing to strip"]}, + request_data={}, + input_type="request", + ) + + assert handler.payloads[0]["texts"] == ["nothing to strip"] + + def test_invalid_regex_fails_at_init(self): + with pytest.raises(ValueError, match="strip_patterns"): + _make_guardrail(_RecordingHandler(), strip_patterns=["(unclosed"]) + + +class TestRequestApplicabilityFilter: + """skip_if_system_prompt_matches / skip_if_first_role_in.""" + + MARKER = "internal-agent-7f3c" + + def _guardrail(self, handler, **options): + return _make_guardrail( + handler, + event_hook=["pre_call", "post_call"], + skip_if_system_prompt_matches=[self.MARKER], + **options, + ) + + @pytest.mark.asyncio + async def test_matching_system_prompt_sends_nothing(self): + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + messages = [ + {"role": "system", "content": f"you are {self.MARKER}"}, + {"role": "user", "content": "hello"}, + ] + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": messages}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(), + ) + + assert handler.calls == [] + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_marker_in_user_message_does_not_skip(self): + """Anchoring to the system message keeps pasted content from skipping the guardrail.""" + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + messages = [ + {"role": "system", "content": "you are helpful"}, + {"role": "user", "content": f"what is {self.MARKER}?"}, + ] + + await guardrail.apply_guardrail( + inputs={"texts": [f"what is {self.MARKER}?"], "structured_messages": messages}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(), + ) + + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_developer_role_prompt_matches(self): + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + messages = [ + {"role": "developer", "content": [{"type": "text", "text": f"id={self.MARKER}"}]}, + {"role": "user", "content": "hello"}, + ] + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": messages}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(), + ) + + assert handler.calls == [] + + @pytest.mark.asyncio + async def test_paired_response_skipped_via_shared_logging_obj(self): + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + logging_obj = _StubLoggingObj() + messages = [{"role": "system", "content": f"you are {self.MARKER}"}] + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": messages}, + request_data={}, + input_type="request", + logging_obj=logging_obj, + ) + await guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + assert handler.calls == [] + + @pytest.mark.asyncio + async def test_paired_response_skipped_via_call_id_cache(self): + """Covers the paths where no logging object reaches the hooks.""" + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + messages = [{"role": "system", "content": f"you are {self.MARKER}"}] + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": messages}, + request_data={"litellm_call_id": "call-xyz"}, + input_type="request", + ) + await guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={"litellm_call_id": "call-xyz"}, + input_type="response", + ) + + assert handler.calls == [] + + @pytest.mark.asyncio + async def test_cache_entry_is_evicted_on_use(self): + """A one-shot entry cannot suppress a later, unrelated response.""" + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + messages = [{"role": "system", "content": f"you are {self.MARKER}"}] + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": messages}, + request_data={"litellm_call_id": "call-xyz"}, + input_type="request", + ) + for _ in range(2): + await guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={"litellm_call_id": "call-xyz"}, + input_type="response", + ) + + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_unmatched_request_leaves_response_scanned(self): + handler = _RecordingHandler() + guardrail = self._guardrail(handler) + logging_obj = _StubLoggingObj() + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": [{"role": "system", "content": "plain"}]}, + request_data={}, + input_type="request", + logging_obj=logging_obj, + ) + await guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + assert [payload["input_type"] for payload in handler.payloads] == ["request", "response"] + + @pytest.mark.asyncio + async def test_skip_decision_does_not_cross_guardrail_instances(self): + handler = _RecordingHandler() + skipping = self._guardrail(handler, name="skipping-guardrail") + other = _make_guardrail( + handler, + name="other-guardrail", + event_hook=["pre_call", "post_call"], + skip_if_system_prompt_matches=["something-else"], + ) + logging_obj = _StubLoggingObj() + messages = [{"role": "system", "content": f"you are {self.MARKER}"}] + + await skipping.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": messages}, + request_data={}, + input_type="request", + logging_obj=logging_obj, + ) + await other.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_skip_if_first_role_in(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, skip_if_first_role_in=["developer"]) + + await guardrail.apply_guardrail( + inputs={ + "texts": ["hello"], + "structured_messages": [ + {"role": "developer", "content": "instructions"}, + {"role": "user", "content": "hello"}, + ], + }, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(), + ) + assert handler.calls == [] + + await guardrail.apply_guardrail( + inputs={ + "texts": ["hello"], + "structured_messages": [{"role": "user", "content": "hello"}], + }, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_id="call-2"), + ) + assert len(handler.calls) == 1 + + +class TestFireAndForget: + """fire_and_forget: dispatch off the request's critical path.""" + + def _guardrail(self, handler, *, max_inflight=10, **options): + dispatcher = BackgroundDispatcher(guardrail_name="test-fire-and-forget", max_inflight=max_inflight) + guardrail = _make_guardrail(handler, dispatcher=dispatcher, fire_and_forget=True, **options) + return guardrail, dispatcher + + @pytest.mark.asyncio + async def test_request_returns_before_the_guardrail_is_called(self): + handler = _RecordingHandler() + guardrail, dispatcher = self._guardrail(handler) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "structured_messages": [{"role": "user", "content": "hello"}]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(), + ) + + # Control is back with the caller while the call is still only scheduled. + assert handler.calls == [] + assert dispatcher.pending_count == 1 + assert result == {"texts": ["hello"], "structured_messages": [{"role": "user", "content": "hello"}]} + + await dispatcher.wait_for_pending() + assert len(handler.calls) == 1 + assert handler.payloads[0]["texts"] == ["hello"] + assert dispatcher.pending_count == 0 + + @pytest.mark.asyncio + async def test_blocked_action_is_ignored(self): + handler = _RecordingHandler(action="BLOCKED") + guardrail, dispatcher = self._guardrail(handler) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + ) + await dispatcher.wait_for_pending() + + assert result == {"texts": ["hello"]} + + @pytest.mark.asyncio + async def test_returned_text_rewrite_is_ignored(self): + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED"]) + guardrail, dispatcher = self._guardrail(handler) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123"]}, + request_data={}, + input_type="request", + ) + await dispatcher.wait_for_pending() + + assert result == {"texts": ["my ssn is 123"]} + + @pytest.mark.asyncio + async def test_failing_endpoint_never_surfaces_to_the_caller(self): + handler = _RecordingHandler(error=httpx.ConnectError("connection refused")) + guardrail, dispatcher = self._guardrail(handler, fail_on_error=True) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + ) + await dispatcher.wait_for_pending() + + assert result == {"texts": ["hello"]} + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_inflight_cap_drops_instead_of_queueing(self): + handler = _RecordingHandler() + guardrail, dispatcher = self._guardrail(handler, max_inflight=1) + + for _ in range(3): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + ) + + assert dispatcher.pending_count == 1 + assert dispatcher.dropped_count == 2 + + await dispatcher.wait_for_pending() + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_payload_shaping_applies_to_the_background_call(self): + handler = _RecordingHandler() + guardrail, dispatcher = self._guardrail(handler, send_images=False, max_text_chars=4) + + await guardrail.apply_guardrail( + inputs={"texts": ["0123456789"], "images": ["data:image/png;base64,AAAA"]}, + request_data={}, + input_type="request", + ) + await dispatcher.wait_for_pending() + + payload = handler.payloads[0] + assert payload["texts"] == ["0123"] + assert "images" not in payload + + def test_streaming_is_forced_to_end_of_stream(self): + """Otherwise every sampled chunk would dispatch its own background call.""" + guardrail, _ = self._guardrail(_RecordingHandler(), streaming_end_of_stream_only=False) + assert guardrail.streaming_end_of_stream_only is True + + def test_max_inflight_must_be_positive(self): + with pytest.raises(ValueError, match="fire_and_forget_max_inflight"): + BackgroundDispatcher(guardrail_name="t", max_inflight=0) + + @pytest.mark.asyncio + async def test_call_type_filter_still_applies(self): + handler = _RecordingHandler() + guardrail, dispatcher = self._guardrail(handler, run_only_on_call_types=["acompletion"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["embed me"]}, + request_data={}, + input_type="request", + logging_obj=_StubLoggingObj(call_type="aembedding"), + ) + await dispatcher.wait_for_pending() + + assert handler.calls == [] + assert dispatcher.pending_count == 0 + + +def _recorded_entries(request_data: dict) -> list: + """The StandardLoggingGuardrailInformation entries the decorator appended.""" + return (request_data.get("metadata") or {}).get("standard_logging_guardrail_information") or [] + + +class TestMaxMessagesBoundsTexts: + """max_messages must bound the flat texts list, not only structured_messages. + + Production translation handlers populate `texts` from every message in the + conversation (one entry per text fragment) alongside `structured_messages`, + so windowing only the structured form would leave payload size proportional + to session length. + """ + + @staticmethod + def _conversation(turns: int) -> tuple[list, list]: + messages = [{"role": "user" if i % 2 == 0 else "assistant", "content": f"turn {i}"} for i in range(turns)] + return [m["content"] for m in messages], messages + + @pytest.mark.asyncio + async def test_texts_are_windowed_with_structured_messages(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, max_messages=2) + texts, messages = self._conversation(6) + + await guardrail.apply_guardrail( + inputs={"texts": texts, "structured_messages": messages}, + request_data={}, + input_type="request", + ) + + payload = handler.payloads[0] + assert payload["texts"] == ["turn 4", "turn 5"] + assert [m["content"] for m in payload["structured_messages"]] == ["turn 4", "turn 5"] + + @pytest.mark.asyncio + async def test_windowed_texts_are_not_written_back(self): + """Windowing shifts positions, so a rewrite cannot be index-mapped.""" + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED-A", "MASKED-B"]) + guardrail = _make_guardrail(handler, max_messages=2) + texts, messages = self._conversation(6) + + result = await guardrail.apply_guardrail( + inputs={"texts": texts, "structured_messages": messages}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == texts + + @pytest.mark.asyncio + async def test_window_larger_than_the_conversation_keeps_write_back(self): + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", texts=["MASKED"]) + guardrail = _make_guardrail(handler, max_messages=10) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["only turn"], "structured_messages": [{"role": "user", "content": "only turn"}]}, + request_data={}, + input_type="request", + ) + + assert handler.payloads[0]["texts"] == ["only turn"] + assert result["texts"] == ["MASKED"] + + +class TestIdentityApplicabilityFilter: + """skip_if_key_alias_in / skip_if_team_id_in: matched on authenticated metadata.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("input_type", ["request", "response"]) + async def test_key_alias_skips_both_hooks_without_correlation(self, input_type): + """Each hook sees the auth metadata, so the response needs no recorded decision.""" + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, skip_if_key_alias_in=["batch-worker"]) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {"user_api_key_alias": "batch-worker"}}, + input_type=input_type, + ) + + assert handler.calls == [] + assert result == {"texts": ["hello"]} + + @pytest.mark.asyncio + async def test_other_key_alias_still_scanned(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, skip_if_key_alias_in=["batch-worker"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {"user_api_key_alias": "prod-app"}}, + input_type="request", + ) + + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_team_id_filter(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, skip_if_team_id_in=["team-internal"]) + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"litellm_metadata": {"user_api_key_team_id": "team-internal"}}, + input_type="response", + ) + assert handler.calls == [] + + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"litellm_metadata": {"user_api_key_team_id": "team-other"}}, + input_type="response", + ) + assert len(handler.calls) == 1 + + @pytest.mark.asyncio + async def test_body_cannot_forge_an_identity_exemption(self): + """A caller putting the alias in its own messages must not be exempted.""" + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, skip_if_key_alias_in=["batch-worker"]) + + await guardrail.apply_guardrail( + inputs={ + "texts": ["batch-worker"], + "structured_messages": [{"role": "system", "content": "batch-worker"}], + }, + request_data={"metadata": {"user_api_key_alias": "prod-app"}}, + input_type="request", + ) + + assert len(handler.calls) == 1 + + +class TestGuardrailInformationScope: + """guardrail_information_scope: how often a logging entry is recorded.""" + + @staticmethod + def _request_data(session_id: str | None = None) -> dict: + data: dict = {"metadata": {}} + if session_id is not None: + data["litellm_session_id"] = session_id + return data + + @pytest.mark.asyncio + async def test_per_call_records_every_invocation(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler) + request_data = self._request_data("session-1") + + for _ in range(3): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert len(_recorded_entries(request_data)) == 3 + + @pytest.mark.asyncio + async def test_per_session_records_only_the_first_call(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, guardrail_information_scope="per_session") + request_data = self._request_data("session-1") + + for _ in range(4): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert len(_recorded_entries(request_data)) == 1 + # The guardrail itself still ran on every call. + assert len(handler.calls) == 4 + + @pytest.mark.asyncio + async def test_per_session_records_once_per_session(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, guardrail_information_scope="per_session") + first = self._request_data("session-1") + second = self._request_data("session-2") + + await guardrail.apply_guardrail(inputs={"texts": ["a"]}, request_data=first, input_type="request") + await guardrail.apply_guardrail(inputs={"texts": ["b"]}, request_data=first, input_type="request") + await guardrail.apply_guardrail(inputs={"texts": ["c"]}, request_data=second, input_type="request") + + assert len(_recorded_entries(first)) == 1 + assert len(_recorded_entries(second)) == 1 + + @pytest.mark.asyncio + async def test_per_session_falls_back_to_per_call_without_a_session_id(self): + """No session id means nothing to dedup against, so entries are not dropped.""" + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, guardrail_information_scope="per_session") + request_data = self._request_data() + + for _ in range(3): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert len(_recorded_entries(request_data)) == 3 + + @pytest.mark.asyncio + async def test_per_session_reads_metadata_session_id(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, guardrail_information_scope="per_session") + request_data = {"metadata": {"session_id": "session-meta"}} + + for _ in range(3): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert len(_recorded_entries(request_data)) == 1 + + @pytest.mark.asyncio + async def test_off_records_nothing_on_success(self): + handler = _RecordingHandler() + guardrail = _make_guardrail(handler, guardrail_information_scope="off") + request_data = self._request_data("session-1") + + for _ in range(3): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert _recorded_entries(request_data) == [] + assert len(handler.calls) == 3 + + @pytest.mark.asyncio + async def test_off_still_records_a_guardrail_failure(self): + """The suppression flag must not swallow the error path.""" + handler = _RecordingHandler(error=httpx.ConnectError("connection refused")) + guardrail = _make_guardrail(handler, guardrail_information_scope="off") + request_data = self._request_data("session-1") + + with pytest.raises(Exception, match="Generic Guardrail API failed"): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert len(_recorded_entries(request_data)) == 1 + + @pytest.mark.asyncio + async def test_per_session_still_records_a_block(self): + handler = _RecordingHandler(action="BLOCKED") + guardrail = _make_guardrail(handler, guardrail_information_scope="per_session") + request_data = self._request_data("session-1") + + for _ in range(2): + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, request_data=request_data, input_type="request" + ) + + assert len(_recorded_entries(request_data)) == 2 + + @pytest.mark.asyncio + async def test_scope_is_per_guardrail_instance(self): + """One guardrail's suppression must not hide another's entry.""" + handler = _RecordingHandler() + suppressed = _make_guardrail(handler, name="quiet", guardrail_information_scope="off") + recording = _make_guardrail(handler, name="loud") + request_data = self._request_data("session-1") + + await suppressed.apply_guardrail(inputs={"texts": ["a"]}, request_data=request_data, input_type="request") + await recording.apply_guardrail(inputs={"texts": ["a"]}, request_data=request_data, input_type="request") + + entries = _recorded_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "loud" + + +class TestConfigValidationWarnings: + """The init-time warnings the config options promise.""" + + def test_unknown_call_type_warns_and_keeps_the_value(self, caplog): + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + guardrail = _make_guardrail(_RecordingHandler(), run_only_on_call_types=["acompletion", "nope"]) + assert "unrecognized call type" in caplog.text + assert guardrail._skip_policy.run_only_on_call_types == frozenset({"acompletion", "nope"}) + + def test_allowlist_and_denylist_together_warns(self, caplog): + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + _make_guardrail( + _RecordingHandler(), + run_only_on_call_types=["acompletion"], + skip_call_types=["aembedding"], + ) + assert "allowlist wins" in caplog.text + + def test_response_only_mode_warns_that_nothing_can_be_skipped(self, caplog): + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + _make_guardrail( + _RecordingHandler(), + event_hook="post_call", + skip_if_system_prompt_matches=["marker"], + ) + assert "request-side hook" in caplog.text + + def test_message_based_filters_warn_about_the_trust_boundary(self, caplog): + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + _make_guardrail(_RecordingHandler(), skip_if_system_prompt_matches=["marker"]) + assert "caller controls" in caplog.text + + def test_identity_filters_do_not_warn(self, caplog): + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + _make_guardrail(_RecordingHandler(), skip_if_key_alias_in=["batch-worker"]) + assert "caller controls" not in caplog.text + + @pytest.mark.asyncio + async def test_excluded_tools_cannot_be_rewritten(self): + handler = _RecordingHandler(action="GUARDRAIL_INTERVENED", tools=[{"type": "function"}]) + guardrail = _make_guardrail(handler, exclude_payload_fields=["tools"]) + original_tools = [{"type": "function", "function": {"name": "run"}}] + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"], "tools": original_tools}, + request_data={}, + input_type="request", + ) + + assert "tools" not in handler.payloads[0] + assert result["tools"] == original_tools diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 3dfb98c12ea..4afe4b1dd16 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -5,6 +5,10 @@ from litellm.proxy.guardrails._content_utils import ( build_inspection_messages, has_non_string_content, iter_message_text, + iter_role_text, + map_content_text, + map_messages_image_urls, + map_messages_text, walk_user_text, ) @@ -563,3 +567,83 @@ def test_build_inspection_messages_custom_tool_call_output(): } msgs = build_inspection_messages(data) assert any("custom-tool-leak" in m["content"] for m in msgs) + + +# ── map_content_text / map_messages_text / map_messages_image_urls ─────────────── + + +def test_map_content_text_transforms_string_and_text_parts_only(): + content = [ + {"type": "text", "text": "keep me"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}}, + "bare fragment", + ] + assert map_content_text(content, lambda t: t.upper()) == [ + {"type": "text", "text": "KEEP ME"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}}, + "BARE FRAGMENT", + ] + + +def test_map_content_text_does_not_mutate_the_input(): + content = [{"type": "text", "text": "original"}] + map_content_text(content, lambda t: "changed") + assert content == [{"type": "text", "text": "original"}] + + +def test_map_messages_text_preserves_every_other_key(): + messages = [ + { + "role": "assistant", + "content": "answer", + "tool_calls": [{"id": "call_1", "function": {"name": "run", "arguments": "{}"}}], + } + ] + mapped = map_messages_text(messages, lambda t: t.upper()) + assert mapped[0]["role"] == "assistant" + assert mapped[0]["content"] == "ANSWER" + assert mapped[0]["tool_calls"] == messages[0]["tool_calls"] + + +def test_map_messages_image_urls_covers_both_part_shapes(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA", "detail": "low"}}, + {"type": "image_url", "image_url": "data:image/png;base64,BBB"}, + {"type": "input_image", "image_url": "data:image/png;base64,CCC"}, + ], + } + ] + mapped = map_messages_image_urls(messages, lambda _: "[omitted]") + content = mapped[0]["content"] + assert content[0] == {"type": "text", "text": "describe"} + assert content[1] == {"type": "image_url", "image_url": {"url": "[omitted]", "detail": "low"}} + assert content[2] == {"type": "image_url", "image_url": "[omitted]"} + assert content[3] == {"type": "input_image", "image_url": "[omitted]"} + + +def test_map_messages_image_urls_leaves_string_content_alone(): + messages = [{"role": "user", "content": "just text"}] + assert map_messages_image_urls(messages, lambda _: "[omitted]") == messages + + +# ── iter_role_text ────────────────────────────────────────────────────────────── + + +def test_iter_role_text_selects_only_the_requested_roles(): + messages = [ + {"role": "system", "content": "system rules"}, + {"role": "developer", "content": [{"type": "text", "text": "developer rules"}]}, + {"role": "user", "content": "user text"}, + ] + assert list(iter_role_text(messages, {"system", "developer"})) == [ + "system rules", + "developer rules", + ] + + +def test_iter_role_text_ignores_non_mapping_entries(): + assert list(iter_role_text(["not a message", {"role": "system", "content": "ok"}], {"system"})) == ["ok"]