From 1d70f3621b96dad95fbf455bdba70c76b6b77f1c Mon Sep 17 00:00:00 2001 From: Cyrill Bannwart Date: Tue, 5 May 2026 20:15:10 +0000 Subject: [PATCH] [Feat] UnifiedLLMGuardrails: streaming action protocol (wait/modify/block) for content rewriting --- .../generic_guardrail_api/__init__.py | 3 + .../generic_guardrail_api.py | 261 +++++---- .../unified_guardrail/unified_guardrail.py | 510 +++++++++++++++++- .../guardrail_hooks/generic_guardrail_api.py | 19 + litellm/types/utils.py | 1 + .../test_unified_guardrail.py | 441 +++++++++++++++ 6 files changed, 1134 insertions(+), 101 deletions(-) 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 bb0d0a99b31..f9a6d76e2cc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -22,6 +22,9 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" litellm_params, "unreachable_fallback", "fail_closed" ), extra_headers=getattr(litellm_params, "extra_headers", None), + supports_streaming_action_protocol=getattr( + litellm_params, "supports_streaming_action_protocol", False + ), guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, 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 790ee31f2e0..6e3277aef8d 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 @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GENERIC_GUARDRAIL_ACTION_BLOCKED, GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, GenericGuardrailAPIResponse, @@ -174,10 +175,14 @@ class GenericGuardrailAPI(CustomGuardrail): And return: { - "action": "BLOCKED" | "NONE" | "GUARDRAIL_INTERVENED", + "action": "BLOCKED" | "NONE" | "GUARDRAIL_INTERVENED" | "WAIT", "blocked_reason": str (optional, only if action is BLOCKED), "text": str (optional, modified text if action is GUARDRAIL_INTERVENED) } + + "WAIT" is only valid in the streaming action protocol (auto-enabled when + the guardrail implements apply_guardrail_action) and only when + is_final=False on the request. """ def __init__( @@ -188,6 +193,7 @@ class GenericGuardrailAPI(CustomGuardrail): additional_provider_specific_params: Optional[Dict[str, Any]] = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", extra_headers: Optional[list] = None, + supports_streaming_action_protocol: bool = False, **kwargs, ): self.async_handler = get_async_httpx_client( @@ -223,6 +229,17 @@ class GenericGuardrailAPI(CustomGuardrail): unreachable_fallback ) + # Streaming action-protocol opt-in. The wrapper class always defines + # apply_guardrail_action (it's just an HTTP call), but whether the + # backing 3rd-party service actually understands `is_final` and the + # WAIT action lives on the service, not on this Python class. Only + # advertise action-protocol support to LiteLLM's auto-detect when + # the operator confirms (via this flag) that their service speaks + # the protocol. Otherwise hide the method so the iterator hook + # falls back to moderation, matching the service's real behavior. + if not supports_streaming_action_protocol: + self.apply_guardrail_action = None # type: ignore[assignment,method-assign] + # Set supported event hooks if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ @@ -367,6 +384,102 @@ class GenericGuardrailAPI(CustomGuardrail): ) raise Exception(f"Generic Guardrail API failed: {str(error)}") + async def _post_guardrail_request( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> Optional[GenericGuardrailAPIResponse]: + """ + Build payload, POST to the guardrail endpoint, parse the response. + + Returns the parsed GenericGuardrailAPIResponse on success, or None + when transport errors are absorbed by the fail-open policy. Does NOT + interpret the action field — callers decide what to do with the response. + """ + # Extract texts and images from inputs + texts = inputs.get("texts", []) + images = inputs.get("images") + tools = inputs.get("tools") + structured_messages = inputs.get("structured_messages") + tool_calls = inputs.get("tool_calls") + model = inputs.get("model") + is_final = inputs.get("is_final") + + if request_data is None: + request_data = {} + + request_body = request_data.get("body") or {} + + additional_params = {**self.additional_provider_specific_params} + dynamic_params = self.get_guardrail_dynamic_request_body_params(request_body) + if dynamic_params: + additional_params.update(dynamic_params) + + user_metadata = 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 = _extract_inbound_headers( + request_data=request_data, + logging_obj=logging_obj, + extra_allowlist=extra_allowlist, + ) + + guardrail_request = GenericGuardrailAPIRequest( + litellm_call_id=logging_obj.litellm_call_id if logging_obj else None, + litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None, + texts=texts, + request_data=user_metadata, + request_headers=inbound_headers, + litellm_version=litellm_version, + images=images, + tools=tools, + structured_messages=structured_messages, + tool_calls=tool_calls, + additional_provider_specific_params=additional_params, + input_type=input_type, + model=model, + is_final=is_final, + ) + + headers = self._build_request_headers() + + try: + response = await self.async_handler.post( + url=self.api_base, + json=guardrail_request.model_dump(mode="json"), + headers=headers, + ) + response.raise_for_status() + response_json = response.json() + + verbose_proxy_logger.debug( + "Generic Guardrail API response: %s", response_json + ) + return GenericGuardrailAPIResponse.from_dict(response_json) + except Timeout as e: + self._handle_guardrail_request_error(e, inputs, input_type, logging_obj) + return None + except httpx.HTTPStatusError as e: + status_code = getattr(getattr(e, "response", None), "status_code", None) + is_unreachable = status_code in (502, 503, 504) + self._handle_guardrail_request_error( + e, inputs, input_type, logging_obj, is_unreachable=is_unreachable + ) + return None + except httpx.RequestError as e: + self._handle_guardrail_request_error(e, inputs, input_type, logging_obj) + return None + except Exception as e: + self._handle_guardrail_request_error( + e, inputs, input_type, logging_obj, is_unreachable=False + ) + return None + @log_guardrail_information async def apply_guardrail( self, @@ -397,117 +510,65 @@ class GenericGuardrailAPI(CustomGuardrail): """ verbose_proxy_logger.debug("Generic Guardrail API: Applying guardrail to text") - # Extract texts and images from inputs texts = inputs.get("texts", []) images = inputs.get("images") tools = inputs.get("tools") - structured_messages = inputs.get("structured_messages") - tool_calls = inputs.get("tool_calls") - model = inputs.get("model") - # Use provided request_data or create an empty dict - if request_data is None: - request_data = {} - - request_body = request_data.get("body") or {} - - # Merge additional provider specific params from config and dynamic params - additional_params = {**self.additional_provider_specific_params} - - # Get dynamic params from request if available - dynamic_params = self.get_guardrail_dynamic_request_body_params(request_body) - if dynamic_params: - additional_params.update(dynamic_params) - - # Extract user API key metadata - user_metadata = 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 = _extract_inbound_headers( + guardrail_response = await self._post_guardrail_request( + inputs=inputs, request_data=request_data, + input_type=input_type, logging_obj=logging_obj, - extra_allowlist=extra_allowlist, ) - # Create request payload - guardrail_request = GenericGuardrailAPIRequest( - litellm_call_id=logging_obj.litellm_call_id if logging_obj else None, - litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None, + # Fail-open transport error path: _handle_guardrail_request_error returned + # a passthrough; mirror its shape here. + if guardrail_response is None: + return_inputs: GenericGuardrailAPIInputs = {} + return_inputs.update(inputs) + return return_inputs + + if guardrail_response.action == GENERIC_GUARDRAIL_ACTION_BLOCKED: + error_message = ( + guardrail_response.blocked_reason or "Content violates policy" + ) + verbose_proxy_logger.warning( + "Generic Guardrail API blocked request: %s", error_message + ) + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message=error_message, + should_wrap_with_default_message=False, + ) + + return self._build_guardrail_return_inputs( texts=texts, - request_data=user_metadata, - request_headers=inbound_headers, - litellm_version=litellm_version, images=images, tools=tools, - structured_messages=structured_messages, - tool_calls=tool_calls, - additional_provider_specific_params=additional_params, - input_type=input_type, - model=model, + guardrail_response=guardrail_response, ) - headers = self._build_request_headers() + async def apply_guardrail_action( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> Optional[GenericGuardrailAPIResponse]: + """ + Streaming action protocol entry point. - try: - # Make the API request - # Use mode="json" to ensure all iterables are converted to lists - response = await self.async_handler.post( - url=self.api_base, - json=guardrail_request.model_dump(mode="json"), - headers=headers, - ) + Returns the raw GenericGuardrailAPIResponse without raising on BLOCKED — + the caller (the unified guardrail iterator hook in action mode) drives + the wait/pass/block/modify state machine. - response.raise_for_status() - response_json = response.json() - - verbose_proxy_logger.debug( - "Generic Guardrail API response: %s", response_json - ) - - guardrail_response = GenericGuardrailAPIResponse.from_dict(response_json) - - # Handle the response - if guardrail_response.action == "BLOCKED": - # Block the request - error_message = ( - guardrail_response.blocked_reason or "Content violates policy" - ) - verbose_proxy_logger.warning( - "Generic Guardrail API blocked request: %s", error_message - ) - raise GuardrailRaisedException( - guardrail_name=GUARDRAIL_NAME, - message=error_message, - should_wrap_with_default_message=False, - ) - - return self._build_guardrail_return_inputs( - texts=texts, - images=images, - tools=tools, - guardrail_response=guardrail_response, - ) - - except GuardrailRaisedException: - raise - except Timeout as e: - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj - ) - except httpx.HTTPStatusError as e: - status_code = getattr(getattr(e, "response", None), "status_code", None) - is_unreachable = status_code in (502, 503, 504) - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj, is_unreachable=is_unreachable - ) - except httpx.RequestError as e: - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj - ) - except Exception as e: - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj, is_unreachable=False - ) + Returns None on fail-open transport error; the caller should treat None + as "no modification, no block, no wait" (i.e. emit accumulated text + unchanged). + """ + return await self._post_guardrail_request( + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=logging_obj, + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index bc46beabc65..f2cbe44b976 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -19,9 +19,23 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings +from litellm.exceptions import GuardrailRaisedException from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GENERIC_GUARDRAIL_ACTION_BLOCKED, + GENERIC_GUARDRAIL_ACTION_GUARDRAIL_INTERVENED, + GENERIC_GUARDRAIL_ACTION_NONE, + GENERIC_GUARDRAIL_ACTION_WAIT, +) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import ( + CallTypes, + CallTypesLiteral, + Delta, + GenericGuardrailAPIInputs, + ModelResponseStream, + StreamingChoices, +) # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message) @@ -29,6 +43,198 @@ A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME = "unified_llm_guardrails" +def _supports_action_protocol(guardrail_to_apply: Any) -> bool: + """ + True iff the guardrail instance advertises support for the streaming + action protocol. + + Auto-detected by looking for a callable `apply_guardrail_action` method + on the guardrail instance. For Python-class-backed guardrails (subclass + `CustomGuardrail` and override the method), presence on the class is the + truth: the method *is* the implementation. + + For HTTP-service-backed wrappers like `GenericGuardrailAPI`, the wrapper + class always *could* call apply_guardrail_action over the wire — but + the protocol's actual semantics (handling `is_final`, returning WAIT) + live on the third-party service. Such wrappers must hide the method on + instances whose backing service doesn't support it (typically by + setting `self.apply_guardrail_action = None` when an operator opt-in + flag is off), so this auto-detection reflects the service's + capabilities rather than the wrapper's. + + Guardrails for which this returns False receive the historical + moderation-mode iterator hook (observe-only); guardrails for which it + returns True drive the wait/pass/block/modify state machine. + """ + return callable(getattr(guardrail_to_apply, "apply_guardrail_action", None)) + + +def _combine_streaming_text(chunks: List[Any]) -> str: + """Concatenate delta.content across all chunks for choice index 0.""" + parts: List[str] = [] + for chunk in chunks: + choices = getattr(chunk, "choices", None) + if not choices: + continue + choice = choices[0] + delta = getattr(choice, "delta", None) + if delta is None: + continue + content = getattr(delta, "content", None) + if content: + parts.append(content) + return "".join(parts) + + +def _build_delta_chunk( + template: Any, + content: str, +) -> ModelResponseStream: + """Build a content-only delta chunk modelled on `template`'s id/model/created.""" + return ModelResponseStream( + id=getattr(template, "id", None), + created=getattr(template, "created", None), + model=getattr(template, "model", None), + object=getattr(template, "object", "chat.completion.chunk"), + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=content, role="assistant"), + finish_reason=None, + ) + ], + ) + + +def _build_terminal_chunk( + template: Any, + finish_reason: str = "stop", +) -> ModelResponseStream: + """Build an empty-content chunk carrying the stream's finish_reason.""" + return ModelResponseStream( + id=getattr(template, "id", None), + created=getattr(template, "created", None), + model=getattr(template, "model", None), + object=getattr(template, "object", "chat.completion.chunk"), + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason=finish_reason, + ) + ], + ) + + +def _last_finish_reason(chunks: List[Any], default: str = "stop") -> str: + """Pull the finish_reason off the most recent chunk that has one.""" + for chunk in reversed(chunks): + choices = getattr(chunk, "choices", None) + if not choices: + continue + fr = getattr(choices[0], "finish_reason", None) + if fr: + return fr + return default + + +def _chunk_tool_call_deltas(chunk: Any) -> List[Any]: + """Return tool_call deltas (if any) for choice 0 of `chunk`, else empty list.""" + choices = getattr(chunk, "choices", None) + if not choices: + return [] + delta = getattr(choices[0], "delta", None) + if delta is None: + return [] + tcs = getattr(delta, "tool_calls", None) + if not tcs: + return [] + return list(tcs) + + +def _accumulate_tool_calls(chunks: List[Any]) -> List[dict]: + """ + Reduce all tool_call deltas across chunks into per-index accumulated state. + + Returns a list of {id, type, function: {name, arguments}} dicts (one per + distinct tool_call index seen), in index order. `arguments` is the + concatenation of all argument deltas seen for that index. + + Used to surface the in-progress tool calls to a guardrail in streaming + action mode so the guardrail can inspect/block based on the args. Note + that mid-stream the args JSON may not yet parse; guardrails that need + a complete payload should return WAIT until they can parse it. + """ + by_index: dict[int, dict] = {} + for chunk in chunks: + for tc in _chunk_tool_call_deltas(chunk): + idx = getattr(tc, "index", None) + if idx is None and isinstance(tc, dict): + idx = tc.get("index") + if idx is None: + continue + slot = by_index.setdefault( + idx, {"id": None, "type": None, "function": {"name": None, "arguments": ""}} + ) + + def _g(obj: Any, key: str) -> Any: + if isinstance(obj, dict): + return obj.get(key) + return getattr(obj, key, None) + + if slot["id"] is None: + slot["id"] = _g(tc, "id") + if slot["type"] is None: + slot["type"] = _g(tc, "type") + fn = _g(tc, "function") + if fn is not None: + if slot["function"]["name"] is None: + name = _g(fn, "name") + if name: + slot["function"]["name"] = name + args = _g(fn, "arguments") + if args: + slot["function"]["arguments"] += args + return [by_index[i] for i in sorted(by_index.keys())] + + +def _replay_tool_call_chunks( + chunks: List[Any], + start_idx: int, + end_idx: int, +) -> List[ModelResponseStream]: + """ + Build emit-ready chunks for tool_call deltas in chunks[start_idx:end_idx]. + + Each output chunk carries only the tool_call portion of the original (no + content, no finish_reason). Chunks with no tool_call deltas in the slice + are skipped. Used in action mode so tool_call deltas reach the client + after the guardrail clears them, even though text content is rebuilt + from the cursor. + """ + out: List[ModelResponseStream] = [] + for chunk in chunks[start_idx:end_idx]: + tcs = _chunk_tool_call_deltas(chunk) + if not tcs: + continue + out.append( + ModelResponseStream( + id=getattr(chunk, "id", None), + created=getattr(chunk, "created", None), + model=getattr(chunk, "model", None), + object=getattr(chunk, "object", "chat.completion.chunk"), + choices=[ + StreamingChoices( + index=0, + delta=Delta(tool_calls=tcs), + finish_reason=None, + ) + ], + ) + ) + return out + + def _get_a2a_request_id( responses_so_far: List[Any], request_data: dict ) -> Optional[str]: @@ -301,6 +507,13 @@ class UnifiedLLMGuardrails(CustomLogger): Supports sampling_rate parameter to control how often chunks are processed. sampling_rate=1 means every chunk, sampling_rate=5 means every 5th chunk, etc. + + Two operating modes, auto-detected per-guardrail: + - moderation (default): observe-only — chunks pass through to client + unmodified; guardrail can block via raised exception. Selected when + the guardrail does not implement apply_guardrail_action. + - action: streaming action protocol (wait/pass/block/modify). Selected + when the guardrail implements apply_guardrail_action. """ global endpoint_guardrail_translation_mappings @@ -360,6 +573,20 @@ class UnifiedLLMGuardrails(CustomLogger): yield item return + # Auto-detect dispatch: guardrails implementing apply_guardrail_action + # opt into the streaming action protocol; everything else stays on + # the historical moderation-mode iterator hook below. + if _supports_action_protocol(guardrail_to_apply): + async for chunk in self._action_mode_stream( + response=response, + guardrail_to_apply=guardrail_to_apply, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + sampling_rate=max(1, int(sampling_rate)), + ): + yield chunk + return + # Initialize translation mappings if needed if endpoint_guardrail_translation_mappings is None: endpoint_guardrail_translation_mappings = ( @@ -520,3 +747,284 @@ class UnifiedLLMGuardrails(CustomLogger): yield error_chunk else: raise + + async def _call_action_guardrail( + self, + *, + guardrail_to_apply: CustomGuardrail, + accumulated_text: str, + chunks: List[Any], + is_final: bool, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + ) -> Optional[Any]: + """ + Call apply_guardrail_action on the guardrail with the accumulated text. + + Returns the raw GenericGuardrailAPIResponse, or None when the guardrail + returns no decision (transport fail-open, or this guardrail integration + does not implement the action protocol). Callers treat None as + "no modification, no block" (NONE-equivalent). + """ + apply_action = getattr(guardrail_to_apply, "apply_guardrail_action", None) + if apply_action is None: + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails: iterator_hook_mode=action requires " + "guardrail.apply_guardrail_action; guardrail=%s does not " + "implement it. Falling back to passthrough.", + getattr(guardrail_to_apply, "guardrail_name", ""), + ) + return None + + inputs: GenericGuardrailAPIInputs = { + "texts": [accumulated_text], + "is_final": is_final, + } + first_chunk = chunks[0] if chunks else None + model = getattr(first_chunk, "model", None) if first_chunk else None + if model: + inputs["model"] = model + + # Surface accumulated tool_calls for inspect-and-block. Modify is not + # supported in action mode — tool_call chunks pass through unchanged + # on emit, regardless of guardrail response. + accumulated_tool_calls = _accumulate_tool_calls(chunks) + if accumulated_tool_calls: + inputs["tool_calls"] = accumulated_tool_calls # type: ignore[typeddict-item] + + return await apply_action( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=request_data.get("litellm_logging_obj"), + ) + + async def _action_mode_stream( # noqa: PLR0915 + self, + *, + response: Any, + guardrail_to_apply: CustomGuardrail, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + sampling_rate: int, + ) -> AsyncGenerator[Any, None]: + """ + Streaming action protocol state machine. + + Buffers upstream chunks, calls the guardrail at sample points (and + on every chunk while in WAIT state), emits delta chunks past the + cursor on NONE/GUARDRAIL_INTERVENED, terminates on BLOCKED, and + always makes a final is_final=True call after upstream EOS. + """ + cursor = 0 + in_wait_state = False + chunk_counter = 0 + all_chunks: List[Any] = [] + # Index into all_chunks of the next chunk whose tool_call deltas (if + # any) have not yet been replayed to the client. Advances on every + # NONE/GUARDRAIL_INTERVENED emit so tool_call deltas reach the client + # in lockstep with text emission, never mid-WAIT. + tool_calls_replayed_through = 0 + template_chunk: Optional[Any] = None + guardrail_name = getattr(guardrail_to_apply, "guardrail_name", GUARDRAIL_NAME) + + async for item in response: + chunk_counter += 1 + all_chunks.append(item) + if template_chunk is None: + template_chunk = item + + sample_due = (chunk_counter % sampling_rate == 0) or in_wait_state + if not sample_due: + continue + + accumulated_text = _combine_streaming_text(all_chunks) + decision = await self._call_action_guardrail( + guardrail_to_apply=guardrail_to_apply, + accumulated_text=accumulated_text, + chunks=all_chunks, + is_final=False, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + ) + + new_text, action = self._resolve_action_decision( + decision=decision, + accumulated_text=accumulated_text, + ) + + if action == GENERIC_GUARDRAIL_ACTION_BLOCKED: + error_message = ( + (decision.blocked_reason if decision is not None else None) + or "Content violates policy" + ) + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails action mode: BLOCKED mid-stream by %s: %s", + guardrail_name, + error_message, + ) + raise GuardrailRaisedException( + guardrail_name=guardrail_name, + message=error_message, + should_wrap_with_default_message=False, + ) + + if action == GENERIC_GUARDRAIL_ACTION_WAIT: + in_wait_state = True + continue + + # NONE or GUARDRAIL_INTERVENED — emit delta past cursor and + # replay any buffered tool_call deltas accumulated since last emit. + self._validate_cursor_monotonic( + new_text=new_text, + cursor=cursor, + guardrail_name=guardrail_name, + is_final=False, + ) + delta = new_text[cursor:] + cursor = len(new_text) + in_wait_state = False + if delta and template_chunk is not None: + yield _build_delta_chunk(template_chunk, delta) + for tc_chunk in _replay_tool_call_chunks( + all_chunks, tool_calls_replayed_through, len(all_chunks) + ): + yield tc_chunk + tool_calls_replayed_through = len(all_chunks) + + # Upstream exhausted — final call with is_final=True. + accumulated_text = _combine_streaming_text(all_chunks) + decision = await self._call_action_guardrail( + guardrail_to_apply=guardrail_to_apply, + accumulated_text=accumulated_text, + chunks=all_chunks, + is_final=True, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + ) + + new_text, action = self._resolve_action_decision( + decision=decision, + accumulated_text=accumulated_text, + ) + + if action == GENERIC_GUARDRAIL_ACTION_BLOCKED: + error_message = ( + (decision.blocked_reason if decision is not None else None) + or "Content violates policy" + ) + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails action mode: BLOCKED at EOS by %s: %s", + guardrail_name, + error_message, + ) + raise GuardrailRaisedException( + guardrail_name=guardrail_name, + message=error_message, + should_wrap_with_default_message=False, + ) + + if action == GENERIC_GUARDRAIL_ACTION_WAIT: + verbose_proxy_logger.error( + "UnifiedLLMGuardrails action mode: %s returned WAIT at " + "is_final=True (action protocol violation)", + guardrail_name, + ) + fallback = getattr( + guardrail_to_apply, "unreachable_fallback", "fail_closed" + ) + if fallback == "fail_open": + new_text = accumulated_text + else: + raise GuardrailRaisedException( + guardrail_name=guardrail_name, + message=( + "guardrail returned WAIT at end of stream " + "(action protocol violation)" + ), + should_wrap_with_default_message=False, + ) + + # At EOS we can't retract bytes already emitted. Soft-fall back to + # the raw accumulated text rather than failing the whole stream. + if len(new_text) < cursor: + verbose_proxy_logger.error( + "UnifiedLLMGuardrails action mode (EOS): %s returned text " + "shorter than already-emitted (cursor=%d, new=%d) — falling " + "back to raw accumulated text", + guardrail_name, + cursor, + len(new_text), + ) + new_text = accumulated_text + + delta = new_text[cursor:] + cursor = len(new_text) + if delta and template_chunk is not None: + yield _build_delta_chunk(template_chunk, delta) + # Replay any tool_call deltas that arrived between the last sample-point + # emit and EOS so the client sees the complete tool-call stream. + for tc_chunk in _replay_tool_call_chunks( + all_chunks, tool_calls_replayed_through, len(all_chunks) + ): + yield tc_chunk + tool_calls_replayed_through = len(all_chunks) + + if template_chunk is not None: + yield _build_terminal_chunk( + template_chunk, + finish_reason=_last_finish_reason(all_chunks), + ) + + @staticmethod + def _resolve_action_decision( + decision: Any, + accumulated_text: str, + ) -> tuple[str, str]: + """Pull (new_text, action) from a GenericGuardrailAPIResponse, with NONE fallback.""" + if decision is None: + return accumulated_text, GENERIC_GUARDRAIL_ACTION_NONE + action = decision.action or GENERIC_GUARDRAIL_ACTION_NONE + if ( + action == GENERIC_GUARDRAIL_ACTION_GUARDRAIL_INTERVENED + and decision.texts + ): + return decision.texts[0], action + if action not in ( + GENERIC_GUARDRAIL_ACTION_NONE, + GENERIC_GUARDRAIL_ACTION_BLOCKED, + GENERIC_GUARDRAIL_ACTION_GUARDRAIL_INTERVENED, + GENERIC_GUARDRAIL_ACTION_WAIT, + ): + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails action mode: unknown action=%r; " + "treating as NONE", + action, + ) + return accumulated_text, GENERIC_GUARDRAIL_ACTION_NONE + return accumulated_text, action + + @staticmethod + def _validate_cursor_monotonic( + new_text: str, + cursor: int, + guardrail_name: str, + is_final: bool, + ) -> None: + if len(new_text) >= cursor: + return + msg = ( + f"guardrail attempted to retract already-emitted content " + f"(cursor={cursor}, returned len={len(new_text)})" + ) + verbose_proxy_logger.error( + "UnifiedLLMGuardrails action mode: %s — %s (is_final=%s)", + guardrail_name, + msg, + is_final, + ) + raise GuardrailRaisedException( + guardrail_name=guardrail_name, + message=f"{msg} — action protocol violation", + should_wrap_with_default_message=False, + ) 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 94f219a5fc6..3eb373c7da9 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -81,6 +81,25 @@ class GenericGuardrailAPIRequest(BaseModel): Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]] ] = None model: Optional[str] = None # the model being used for the LLM call + is_final: Optional[bool] = Field( + default=None, + description=( + "Streaming action protocol: True iff this is the final guardrail call " + "for the stream (no more upstream chunks will arrive). The guardrail " + "MUST NOT return action=WAIT when is_final=True. Unset/None outside " + "the streaming action-protocol path." + ), + ) + + +# Action values returned by Generic Guardrail API responses. +# WAIT is used only by the streaming action protocol (auto-enabled when +# the guardrail implements apply_guardrail_action) and only when +# is_final=False on the request. +GENERIC_GUARDRAIL_ACTION_NONE = "NONE" +GENERIC_GUARDRAIL_ACTION_BLOCKED = "BLOCKED" +GENERIC_GUARDRAIL_ACTION_GUARDRAIL_INTERVENED = "GUARDRAIL_INTERVENED" +GENERIC_GUARDRAIL_ACTION_WAIT = "WAIT" class GenericGuardrailAPIResponse: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c05c46e0d45..68ef08e002b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3673,3 +3673,4 @@ class GenericGuardrailAPIInputs(TypedDict, total=False): AllMessageValues ] # structured messages sent to the LLM - indicates if text is from system or user model: Optional[str] # the model being used for the LLM call + is_final: bool # streaming action protocol: True iff this is the final guardrail call for the stream diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 2418d7af04b..c7c31938f5b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -351,6 +351,447 @@ class TestUnifiedLLMGuardrails: f"Expected non-empty content for every streamed chunk." ) + class TestActionMode: + """Streaming action protocol (auto-enabled when the guardrail + implements apply_guardrail_action).""" + + @staticmethod + def _make_chunk(content=None, tool_calls=None, finish_reason=None): + return ModelResponseStream( + id="cmpl-test", + created=0, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + tool_calls=tool_calls, + ), + finish_reason=finish_reason, + ) + ], + ) + + @staticmethod + def _content_chunks(parts, finish_reason="stop"): + n = len(parts) + + async def gen(): + for i, p in enumerate(parts): + yield TestUnifiedLLMGuardrails.TestActionMode._make_chunk( + content=p, + finish_reason=finish_reason if i == n - 1 else None, + ) + + return gen() + + @staticmethod + async def _collect(generator): + out = [] + async for item in generator: + out.append(item) + return out + + class _ScriptedActionGuardrail(CustomGuardrail): + """Drives apply_guardrail_action with a scripted sequence of decisions.""" + + def __init__(self, scripted_decisions, sampling_rate=2): + super().__init__(guardrail_name="scripted-action-guardrail") + # Each decision: (action, texts, blocked_reason). + # apply_guardrail_action's presence on this class is what the + # iterator hook auto-detects to drive action mode. + self.script = list(scripted_decisions) + self.calls = [] + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = False + self.unreachable_fallback = "fail_closed" + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail_action( + self, *, inputs, request_data, input_type, logging_obj=None + ): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + if not self.script: + raise AssertionError( + "scripted guardrail ran out of decisions" + ) + action, texts, blocked_reason = self.script.pop(0) + self.calls.append( + { + "text": (inputs.get("texts") or [""])[0], + "is_final": inputs.get("is_final"), + "tool_calls": inputs.get("tool_calls"), + "action": action, + } + ) + return GenericGuardrailAPIResponse( + action=action, + texts=[texts] if texts is not None else None, + blocked_reason=blocked_reason, + ) + + @pytest.mark.asyncio + async def test_auto_detect_falls_back_to_moderation(self): + """A guardrail without apply_guardrail_action gets moderation mode. + + The iterator hook dispatches on the presence of a callable + apply_guardrail_action method — RecordingGuardrail doesn't have + one, so the historical observe-only iterator hook should run and + yield the original chunks unmodified. + """ + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + assert not callable( + getattr(guardrail, "apply_guardrail_action", None) + ), "this test relies on RecordingGuardrail not implementing the action protocol" + + chunks = [ + ModelResponseStream( + choices=[ + StreamingChoices( + delta=Delta(content=f"w{i}", role="assistant"), + finish_reason=None, + ) + ], + ) + for i in range(3) + ] + + async def upstream(): + for c in chunks: + yield c + + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4"} + + out = await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=upstream(), + request_data=request_data, + ) + ) + # Moderation passes through original chunks unchanged. + assert len(out) == 3 + assert "".join( + c.choices[0].delta.content for c in out if c.choices[0].delta.content + ) == "w0w1w2" + + @pytest.mark.asyncio + async def test_action_mode_emits_modified_text(self): + """GUARDRAIL_INTERVENED: client receives modified text only.""" + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [ + ("GUARDRAIL_INTERVENED", "ABCD", None), + ("GUARDRAIL_INTERVENED", "ABCDEFGH", None), + ("GUARDRAIL_INTERVENED", "ABCDEFGH", None), + ], + sampling_rate=2, + ) + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + request_data = {"guardrail_to_apply": guardrail} + + out = await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks( + ["ab", "cd", "ef", "gh"] + ), + request_data=request_data, + ) + ) + text = "".join( + c.choices[0].delta.content or "" for c in out + ) + assert text == "ABCDEFGH" + assert out[-1].choices[0].finish_reason == "stop" + assert [c["is_final"] for c in guardrail.calls] == [False, False, True] + + @pytest.mark.asyncio + async def test_action_mode_buffers_during_wait(self): + """WAIT yields nothing; on resume, the full delta reaches the client.""" + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [ + ("WAIT", None, None), + ("WAIT", None, None), + ("WAIT", None, None), + ("GUARDRAIL_INTERVENED", "FINAL", None), + ], + sampling_rate=2, + ) + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + out = await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks( + ["ab", "cd", "ef", "gh"] + ), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + text = "".join(c.choices[0].delta.content or "" for c in out) + assert text == "FINAL" + # WAIT collapses sample-rate to 1: every subsequent chunk triggers a call. + assert [c["is_final"] for c in guardrail.calls] == [ + False, + False, + False, + True, + ] + + @pytest.mark.asyncio + async def test_action_mode_blocks_mid_stream(self): + """BLOCKED raises GuardrailRaisedException with the blocked_reason.""" + from litellm.exceptions import GuardrailRaisedException + + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [("BLOCKED", None, "policy denied")], + sampling_rate=2, + ) + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + with pytest.raises(GuardrailRaisedException, match="policy denied"): + await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks( + ["ab", "cd"] + ), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + + @pytest.mark.asyncio + async def test_action_mode_wait_at_eos_is_violation(self): + """WAIT at is_final=True is a protocol violation; fail_closed raises.""" + from litellm.exceptions import GuardrailRaisedException + + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + # sampling_rate=100 ensures only the EOS call is made + [("WAIT", None, None)], + sampling_rate=100, + ) + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + with pytest.raises( + GuardrailRaisedException, match="end of stream" + ): + await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks( + ["ab"] + ), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + + @pytest.mark.asyncio + async def test_action_mode_modify_shrink_is_violation(self): + """GUARDRAIL_INTERVENED with text shorter than cursor raises.""" + from litellm.exceptions import GuardrailRaisedException + + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [ + ("GUARDRAIL_INTERVENED", "abcd", None), # cursor=4 + ("GUARDRAIL_INTERVENED", "ab", None), # shrink → violation + ], + sampling_rate=2, + ) + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + with pytest.raises( + GuardrailRaisedException, match="retract already-emitted" + ): + await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks( + ["ab", "cd", "ef", "gh"] + ), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + + @pytest.mark.asyncio + async def test_action_mode_chunk_straddling_surrogate(self): + """Surrogate token spans chunks: WAIT until complete, then INTERVENED.""" + handler = UnifiedLLMGuardrails() + # Upstream produces "before [EMA" / "IL_1] after" — surrogate split. + # Guardrail returns WAIT on partial, GUARDRAIL_INTERVENED on complete. + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [ + ("WAIT", None, None), + ( + "GUARDRAIL_INTERVENED", + "before [EMAIL_1] after", + None, + ), + ( + "GUARDRAIL_INTERVENED", + "before [EMAIL_1] after", + None, + ), + ], + sampling_rate=1, + ) + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + out = await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=TestUnifiedLLMGuardrails.TestActionMode._content_chunks( + ["before [EMA", "IL_1] after"] + ), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + text = "".join(c.choices[0].delta.content or "" for c in out) + # Client never sees the leaky partial; full surrogate emitted as one delta. + assert text == "before [EMAIL_1] after" + assert "[EMA" not in "".join( + c.choices[0].delta.content + for c in out + if c.choices[0].delta.content + and len(c.choices[0].delta.content) < 10 + ), "no partial-surrogate sub-emit" + + @pytest.mark.asyncio + async def test_action_mode_tool_calls_surfaced_to_guardrail(self): + """Accumulated tool_calls reach the guardrail via inputs.tool_calls.""" + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [ + ("NONE", None, None), + ("NONE", None, None), + ("NONE", None, None), + ], + sampling_rate=2, + ) + + async def upstream(): + _make = TestUnifiedLLMGuardrails.TestActionMode._make_chunk + yield _make(content="Hi! ") + yield _make( + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"loc":"', + }, + } + ] + ) + yield _make( + tool_calls=[ + { + "index": 0, + "function": {"arguments": 'NYC"}'}, + } + ] + ) + yield _make(content="Done.", finish_reason="stop") + + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + out = await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=upstream(), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + # Final guardrail call sees the complete accumulated tool_call. + final_call = guardrail.calls[-1] + assert final_call["is_final"] is True + assert final_call["tool_calls"] is not None + assert len(final_call["tool_calls"]) == 1 + tc = final_call["tool_calls"][0] + assert tc["id"] == "call_1" + assert tc["function"]["name"] == "get_weather" + assert tc["function"]["arguments"] == '{"loc":"NYC"}' + + # Client sees both tool_call delta chunks (replayed at emit / EOS). + tc_chunks = [ + c + for c in out + if c.choices[0].delta and c.choices[0].delta.tool_calls + ] + assert len(tc_chunks) == 2 + + # Text content reaches client correctly. + text = "".join(c.choices[0].delta.content or "" for c in out) + assert text == "Hi! Done." + + @pytest.mark.asyncio + async def test_action_mode_tool_calls_blockable(self): + """BLOCKED based on tool_call args terminates the stream.""" + from litellm.exceptions import GuardrailRaisedException + + handler = UnifiedLLMGuardrails() + guardrail = TestUnifiedLLMGuardrails.TestActionMode._ScriptedActionGuardrail( + [("BLOCKED", None, "denied tool_call")], + sampling_rate=2, + ) + + async def upstream(): + _make = TestUnifiedLLMGuardrails.TestActionMode._make_chunk + yield _make(content="Hi") + yield _make( + tool_calls=[ + { + "index": 0, + "id": "x", + "function": { + "name": "rm", + "arguments": '{"path":"/"}', + }, + } + ], + finish_reason="tool_calls", + ) + + user = UserAPIKeyAuth( + api_key="k", request_route="/v1/chat/completions" + ) + with pytest.raises( + GuardrailRaisedException, match="denied tool_call" + ): + await TestUnifiedLLMGuardrails.TestActionMode._collect( + handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user, + response=upstream(), + request_data={"guardrail_to_apply": guardrail}, + ) + ) + class TestOCRGuardrailE2E: """End-to-end tests: UnifiedLLMGuardrails -> OCRHandler."""