mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
[Feat] UnifiedLLMGuardrails: streaming action protocol (wait/modify/block) for content rewriting
This commit is contained in:
Cyrill Bannwart
parent
f318ef03bd
commit
1d70f3621b
6 changed files with 1134 additions and 101 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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", "<unknown>"),
|
||||
)
|
||||
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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue