scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #38721 from BerriAI/litellm_fix_post_call_policy_pipeline

Browse source 
fix(policy_engine): execute post_call guardrail pipelines on responses and streams
This commit is contained in:
Mateo Wang 2026-09-08 15:19:19 -07:00 • committed by GitHub
commit 047b8bef31
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 2913 additions and 145 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

83
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -13,9 +13,10 @@ Pattern Overview:
"""
import json
from collections.abc import Mapping, Sequence
from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from itertools import chain, repeat
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import ReadOnly, TypedDict, assert_never
@ -29,6 +30,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
anthropic_tool_name,
@ -168,6 +170,8 @@ class AnthropicMessagesHandler(BaseTranslation):
them through guardrail rewrites; downstream provider handling is out of scope.
"""
delivers_ended_stream_text_rewrites = True
def __init__(self):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
@ -1014,11 +1018,17 @@ class AnthropicMessagesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Sequence[object]:
"""
Process output streaming response by applying guardrails to text content.
Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far.
With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite
written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked);
a rewrite on a stream that never reported a ``stop_reason`` has no write-back and is reported as
undeliverable, so the pipeline executor discards it and releases the original chunks.
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -1065,6 +1075,15 @@ class AnthropicMessagesHandler(BaseTranslation):
responses_so_far, request_data
)
raise
guardrailed_texts: Final = _guardrailed_inputs.get("texts")
if (
deliver_ended_stream_rewrites
and isinstance(string_so_far, str)
and string_so_far
and guardrailed_texts
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
return responses_so_far
@ -1087,6 +1106,11 @@ class AnthropicMessagesHandler(BaseTranslation):
if e.original_response is None:
e.original_response = self._build_streaming_usage_response(responses_so_far, request_data)
raise
unended_texts: Final = _guardrailed_inputs.get("texts")
if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
return responses_so_far
def _prepare_request_data(
@ -1180,6 +1204,63 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
@staticmethod
def _write_ended_stream_text_rewrite(
responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: str,
) -> None:
"""Deliver an ended-stream guardrail text rewrite by rewriting the
buffered chunks in place: the first ``text_delta`` carries the full
rewritten text and every later one is blanked, leaving the surrounding
message and content-block framing untouched. Handles both chunk formats
this stream carries (parsed event dicts and raw SSE bytes)."""
replacements: Final = chain((rewritten_text,), repeat(""))
for idx, item in enumerate(responses_so_far):
if isinstance(item, dict):
delta = item.get("delta")
if item.get("type") == "content_block_delta" and isinstance(delta, dict):
if delta.get("type") == "text_delta":
delta["text"] = next(replacements)
elif isinstance(item, (bytes, bytearray)):
responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer
AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
)
@staticmethod
def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
"""Rewrite every ``text_delta`` data line in one SSE chunk with the next
replacement text, leaving all other events and framing byte-identical."""
try:
decoded: Final = sse_bytes.decode("utf-8")
except UnicodeDecodeError:
return sse_bytes
return "\n\n".join(
AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
).encode("utf-8")
@staticmethod
def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
@staticmethod
def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
if not line.startswith("data:"):
return line
try:
data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads(
line[len("data:") :].strip()
)
except json.JSONDecodeError:
return line
if not isinstance(data, dict) or data.get("type") != "content_block_delta":
return line
delta: Final = data.get("delta")
if not isinstance(delta, dict) or delta.get("type") != "text_delta":
return line
return "data: " + json.dumps(
{**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts
)
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
return StreamingScanKey(

16
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, Optional
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
if TYPE_CHECKING:
from fastapi import HTTPException
@ -52,6 +52,14 @@ class StreamingScanKey:
class BaseTranslation(ABC):
delivers_ended_stream_text_rewrites: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` accepts
``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
stream, writes guardrail text rewrites back across ``responses_so_far`` so
a buffered pipeline can release rewritten chunks. Tool-call rewrites, and
text rewrites on every other translation, are undeliverable: the pipeline
executor discards them and releases the original chunks."""
@staticmethod
def transform_user_api_key_dict_to_metadata(
user_api_key_dict: Any | None,
@ -157,6 +165,7 @@ class BaseTranslation(ABC):
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Any:
"""
Process output streaming response with guardrails.
@ -164,6 +173,11 @@ class BaseTranslation(ABC):
Optional to override in subclasses. ``stream_transform_sink`` is the
out-parameter used by handlers that support streaming text
transformations (see ``StreamTransformSink``); base handlers ignore it.
``deliver_ended_stream_rewrites`` is passed True only when the caller
holds the whole buffered stream and the subclass declares
``delivers_ended_stream_text_rewrites``: the handler then writes
guardrail text rewrites back across ``responses_so_far`` instead of
discarding them.
"""
return responses_so_far

126
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -78,6 +78,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
Convert chat completions request data to OpenAI-spec structured messages.
@ -453,6 +455,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
"""
Process output streaming responses by applying guardrails to text content.
@ -467,6 +470,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
accumulated text (``responses_so_far`` is left untouched so it stays
a correct raw accumulator across rounds) and the guardrailed text
plus requested holdback are reported per choice on the sink.
deliver_ended_stream_rewrites: When True and the buffered stream has
ended, guardrail text rewrites are written back across
``responses_so_far`` (full rewritten text in each choice's first
content-carrying chunk, the rest blanked) instead of discarded.
Returns:
The (unmodified) list of responses.
@ -492,6 +499,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
)
async def _process_streaming_block_only(
@ -502,27 +510,23 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict | None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
terminate the stream. Text rewrites are not propagated to the client here
(see ``_process_streaming_transform`` for the incremental_diff path)."""
(see ``_process_streaming_transform`` for the incremental_diff path) unless
``deliver_ended_stream_rewrites`` opts the ended-stream branch in."""
has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
if has_stream_ended:
# convert to model response
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
# run process_output_response
await self.process_output_response(
response=model_response,
await self._process_ended_stream(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
)
return responses_so_far
# Step 0: Check if any response has text content to process
@ -595,6 +599,39 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return responses_so_far
async def _process_ended_stream(
self,
*,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict[str, object] | None, # mutable-ok: same request-payload shape the hooks take
deliver_ended_stream_rewrites: bool,
) -> None:
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text rewrite
back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
await self.process_output_response(
response=model_response,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if deliver_ended_stream_rewrites:
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
def build_stream_error_items(
self,
exc: "HTTPException",
@ -745,8 +782,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""
combined_texts: Final[dict[tuple[int, int | None], str]] = {}
for response_idx, response in enumerate(responses_so_far):
for choice_idx, choice in enumerate(response.choices):
for response in responses_so_far:
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
elif isinstance(choice, litellm.Choices):
@ -759,7 +796,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(content, str):
# String content - accumulate for this choice
str_key: tuple[int, int | None] = (choice_idx, None)
str_key: tuple[int, int | None] = (choice.index, None)
if str_key not in combined_texts:
combined_texts[str_key] = ""
combined_texts[str_key] += content
@ -770,7 +807,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_str = content_item.get("text")
if text_str:
list_key: tuple[int, int | None] = (
choice_idx,
choice.index,
content_idx,
)
if list_key not in combined_texts:
@ -960,6 +997,52 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "name" in func_dict:
existing_tool_call.function.name = func_dict["name"]
@staticmethod
def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]:
return tuple(
choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices
)
async def _write_ended_stream_text_rewrites(
self,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrailed_response: "ModelResponse",
pre_guardrail_texts: tuple[str | None, ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail text rewrites back across the buffered
chunks: the full rewritten text lands in the choice's first
content-carrying chunk and the rest are blanked, the same shape the
in-flight write-back uses. Chunks carrying only finish_reason or usage
stay untouched. A rewrite on a stream carrying more than one distinct
choice index is reported as undeliverable, so the pipeline executor
discards it and releases the original chunks."""
post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response)
changed: Final = tuple(
after
for before, after in zip(pre_guardrail_texts, post_guardrail_texts)
if before is not None and after is not None and after != before
)
if not changed:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
if len(stream_choice_indices) != 1:
# stream_chunk_builder collapses every choice into one index-0
# choice, so a rewrite of the rebuilt response cannot be attributed
# back to a single choice on an n>1 stream: report it undeliverable
# rather than deliver the rewrite on the wrong choice
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
target_choice_index: Final = next(iter(stream_choice_indices))
await self._apply_guardrail_responses_to_output_streaming(
responses=responses_so_far,
guardrailed_texts=list(changed), # mutable-ok: callee takes lists
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
)
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],
@ -975,7 +1058,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Args:
responses: List of ModelResponseStream objects to modify
guardrailed_texts: List of guardrailed text responses (combined from all chunks)
task_mappings: List of tuples (choice_idx, content_idx)
task_mappings: List of tuples (choice_idx, content_idx), where choice_idx
is the choice's ``index`` field, not its position in a chunk's list
Override this method to customize how responses are applied to streaming responses.
"""
@ -991,9 +1075,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# Key: (choice_idx, content_idx), Value: boolean (True if already set)
already_set: Final[dict[tuple[int, int | None], bool]] = {}
# Iterate through all responses and update content
for response_idx, response in enumerate(responses):
for choice_idx_in_response, choice in enumerate(response.choices):
# Iterate through all responses and update content, matching each chunk's
# choice by its index field: on n>1 streams a chunk usually carries one
# choice at list position 0 whose index names the logical choice.
for response in responses:
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
elif isinstance(choice, litellm.Choices):
@ -1006,7 +1092,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(content, str):
# String content
str_key: tuple[int, int | None] = (choice_idx_in_response, None)
str_key: tuple[int, int | None] = (choice.index, None)
if str_key in guardrail_map:
if str_key not in already_set:
# First chunk - set the complete guardrailed text
@ -1027,7 +1113,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
for content_idx, content_item in enumerate(content):
if "text" in content_item:
list_key: tuple[int, int | None] = (
choice_idx_in_response,
choice.index,
content_idx,
)
if list_key in guardrail_map:

145
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -33,7 +33,7 @@ import time
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from itertools import accumulate, chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
@ -49,6 +49,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_stream_usage,
@ -118,6 +119,15 @@ class ResponsesStreamChunk(TypedDict, total=False):
content_index: ReadOnly[int]
_TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
{
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
}
)
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
)
@ -330,6 +340,8 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
Convert Responses API request data to OpenAI-spec structured messages.
@ -667,6 +679,8 @@ class OpenAIResponsesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list[Any]:
"""
Process output streaming response by applying guardrails to text content.
@ -675,10 +689,18 @@ class OpenAIResponsesHandler(BaseTranslation):
chunk, apply the guardrail, then write the result back in-place so the
caller sees the modified content (e.g. PII tokens replaced).
For ``response.completed`` events (the normal end-of-stream signal) we
use the same per-item extraction + task-mapping approach as
``process_output_response`` so that unmasking / blocking works correctly
for every output item.
For terminal envelope events (``response.completed``, and equally
``response.incomplete`` / ``response.failed``, whose envelopes carry the
partial output) we use the same per-item extraction + task-mapping
approach as ``process_output_response`` so that unmasking / blocking
works correctly for every output item. With
``deliver_ended_stream_rewrites`` the earlier text-carrying events
(``response.output_text.delta`` / ``.done``,
``response.content_part.done``, ``response.output_item.done``) are synced
to the rewritten envelope too, so a client reading deltas sees the
rewrite instead of the raw model output; a rewrite observed where no
write-back is possible is reported as undeliverable, so the pipeline
executor discards it and releases the original events.
"""
if not responses_so_far:
return responses_so_far
@ -690,14 +712,16 @@ class OpenAIResponsesHandler(BaseTranslation):
return responses_so_far
# ------------------------------------------------------------------ #
# Case 1: response.completed — full response is available in the #
# final chunk; iterate output items, apply guardrail, write back. #
# Case 1: terminal envelope events (completed/incomplete/failed). #
# the accumulated response is available in the final chunk; iterate #
# output items, apply guardrail, write back. Falls through to the #
# string fallback when the envelope yields nothing to check. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.completed":
if final_chunk.get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES:
response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {}
if not hasattr(response_obj, "get"):
return responses_so_far
outputs: Final[Sequence[object]] = response_obj.get("output") or []
outputs: Final[Sequence[object]] = (
(response_obj.get("output") or []) if hasattr(response_obj, "get") else []
)
texts_to_check: Final[list[str]] = []
tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = []
@ -747,11 +771,25 @@ class OpenAIResponsesHandler(BaseTranslation):
responses=guardrailed_texts,
task_mappings=task_mappings,
)
return responses_so_far
if deliver_ended_stream_rewrites:
rewrites_by_position: Final = MappingProxyType(
{
task_mappings[task_idx]: rewritten
for task_idx, rewritten in enumerate(guardrailed_texts)
if task_idx < len(texts_to_check) and rewritten != texts_to_check[task_idx]
}
)
if rewrites_by_position:
self._sync_stream_events_with_rewrites(
stream_events=responses_so_far[:-1],
rewrites_by_position=rewrites_by_position,
)
return responses_so_far
# ------------------------------------------------------------------ #
# Case 2: response.output_item.done — extract tool calls only. #
# Case 2: response.output_item.done — extract tool calls only, then #
# fall through to the text fallback when a caller expects rewrites #
# delivered, so a truncated buffer still reports text undeliverable. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.output_item.done":
model_response_stream: Final = (
@ -769,12 +807,14 @@ class OpenAIResponsesHandler(BaseTranslation):
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
if not deliver_ended_stream_rewrites:
return responses_so_far
# ------------------------------------------------------------------ #
# Fallback: apply guardrail to the accumulated text string. #
# No structured write-back is possible here; guardrails that only #
# need to block/flag (not rewrite) still work correctly. #
# need to block/flag (not rewrite) still work correctly, and a #
# rewrite a caller expects delivered is reported undeliverable. #
# ------------------------------------------------------------------ #
string_so_far: Final = self.get_streaming_string_so_far(responses_so_far)
if string_so_far:
@ -784,28 +824,83 @@ class OpenAIResponsesHandler(BaseTranslation):
)
if response_model:
fallback_inputs["model"] = response_model
await guardrail_to_apply.apply_guardrail(
fallback_outputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=fallback_inputs,
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,
)
fallback_texts: Final = fallback_outputs.get("texts")
if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
return responses_so_far
@staticmethod
def _write_event_field(event: object, field: str, value: str) -> None:
if isinstance(event, dict):
event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place
else:
setattr(event, field, value)
def _sync_stream_events_with_rewrites(
self,
stream_events: Sequence[Any],
rewrites_by_position: Mapping[tuple[int, int], str],
) -> None:
"""Sync pre-completion stream events with the rewritten completed
response, keyed by ``(output_index, content_index)``: the first
``output_text.delta`` for a rewritten item carries the full rewritten
text and the rest are blanked, while ``output_text.done``,
``content_part.done``, and ``output_item.done`` events carry the full
rewritten text, so every event a client may read agrees with the
rewritten ``response.completed`` payload."""
delta_replacements: Final = MappingProxyType(
{position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()}
)
for event in stream_events:
if not (isinstance(event, dict) or hasattr(event, "get")):
continue
event_type = event.get("type")
output_index = event.get("output_index")
content_index = event.get("content_index")
if event_type == "response.output_item.done" and isinstance(output_index, int):
self._sync_output_item_done_event(event.get("item"), output_index, rewrites_by_position)
continue
if not isinstance(output_index, int) or not isinstance(content_index, int):
continue
position = (output_index, content_index)
if event_type == "response.output_text.delta" and position in delta_replacements:
self._write_event_field(event, "delta", next(delta_replacements[position]))
elif event_type == "response.output_text.done" and position in rewrites_by_position:
self._write_event_field(event, "text", rewrites_by_position[position])
elif event_type == "response.content_part.done" and position in rewrites_by_position:
part = event.get("part")
if isinstance(part, dict) or hasattr(part, "text"):
self._write_event_field(part, "text", rewrites_by_position[position])
@staticmethod
def _sync_output_item_done_event(
item: object,
output_index: int,
rewrites_by_position: Mapping[tuple[int, int], str],
) -> None:
content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None)
if not isinstance(content, list):
return
for (item_idx, content_idx), rewritten in rewrites_by_position.items():
if item_idx != output_index or content_idx >= len(content):
continue
OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if the streaming has ended.
"""
if not responses_so_far:
return False
terminal_types: Final = frozenset(
(
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
)
)
return stream_item_field(responses_so_far[-1], "type") in terminal_types
return stream_item_field(responses_so_far[-1], "type") in _TERMINAL_ENVELOPE_EVENT_TYPES
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
if not responses_so_far or not hasattr(responses_so_far[-1], "get"):

69
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -19,7 +19,7 @@ from litellm.cost_calculator import _infer_call_type
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.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
@ -69,6 +69,36 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran
return translation
def resolve_endpoint_translation(
user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None
) -> "tuple[str, BaseTranslation] | None":
"""
Resolve the endpoint guardrail translation for a streamed response: the
request route wins, falling back to inferring the call type from the first
response chunk (the same resolution order the streaming iterator hook uses).
Returns None when the call type is unresolvable or has no translation.
"""
route_call_types: Final = (
get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None
)
call_type: Final = (
route_call_types[0].value
if route_call_types
else (
_infer_call_type(call_type=None, completion_response=first_response_item)
if first_response_item is not None
else None
)
)
if call_type is None:
return None
try:
handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type))
except ValueError:
return None
return call_type, handler_cls()
def _chunk_choices(item: object) -> Sequence[object]:
choices: Final[Sequence[object]] = getattr(item, "choices", None) or []
return choices
@ -343,7 +373,7 @@ class UnifiedLLMGuardrails(CustomLogger):
return response
async def _handle_streaming_block(
async def handle_streaming_block(
self,
exc: "ModifyResponseException",
endpoint_translation: _EndpointTranslation,
@ -399,7 +429,7 @@ class UnifiedLLMGuardrails(CustomLogger):
return None
return call_type
async def _emit_streaming_http_error(
async def emit_streaming_http_error(
self,
exc: HTTPException,
call_type: str | None,
@ -592,7 +622,7 @@ class UnifiedLLMGuardrails(CustomLogger):
except ModifyResponseException as e:
if e.original_response is None:
e.original_response = responses_so_far
async for block_chunk in self._handle_streaming_block(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=bool(responses_yielded),
@ -601,7 +631,7 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
raise _StreamTerminated()
except HTTPException as e:
async for error_item in self._emit_streaming_http_error(
async for error_item in self.emit_streaming_http_error(
e,
call_type,
responses_so_far,
@ -781,7 +811,7 @@ class UnifiedLLMGuardrails(CustomLogger):
except ModifyResponseException as e:
if e.original_response is None:
e.original_response = responses_so_far
async for block_chunk in self._handle_streaming_block(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=bool(responses_yielded),
@ -869,6 +899,14 @@ class UnifiedLLMGuardrails(CustomLogger):
choices: Final = _chunk_choices(item)
return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object:
"""Streaming flag resolution order (later wins): default < guardrail
attribute < guardrail_config dict < this callback's optional_params."""
attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default)
config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None)
config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value
return self.optional_params.get(name, config_value)
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -897,17 +935,8 @@ class UnifiedLLMGuardrails(CustomLogger):
if guardrail_to_apply is None:
guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
# Get streaming configuration. Resolution order (later wins): default
# < guardrail attribute < guardrail_config dict < this callback's
# optional_params.
def _streaming_flag(name: str, default: object) -> Any:
value = default
if guardrail_to_apply is not None:
value = getattr(guardrail_to_apply, name, value)
config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {})
if isinstance(config, dict):
value = config.get(name, value)
return self.optional_params.get(name, value)
return self.resolve_streaming_flag(guardrail_to_apply, name, default)
sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
# Only apply the guardrail at end of stream (not per chunk).
@ -1091,7 +1120,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# The current chunk was appended to responses_so_far but not
# yet yielded, so exclude it: the continuation must reflect
# only what the client has actually received.
async for block_chunk in self._handle_streaming_block(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=chunks_yielded,
@ -1101,7 +1130,7 @@ class UnifiedLLMGuardrails(CustomLogger):
return
except HTTPException as e:
# Response already started (we already yielded chunks); cannot send 400.
async for error_item in self._emit_streaming_http_error(
async for error_item in self.emit_streaming_http_error(
e,
call_type,
responses_so_far,
@ -1175,7 +1204,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# terminating SSE sequence with the block message rather than
# propagating into a bare error blob that truncates the stream.
# The withheld original chunks are never released.
async for block_chunk in self._handle_streaming_block(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=bool(responses_yielded),
@ -1184,7 +1213,7 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
return
except HTTPException as e:
async for error_item in self._emit_streaming_http_error(
async for error_item in self.emit_streaming_http_error(
e,
call_type,
responses_so_far,

7
litellm/proxy/litellm_pre_call_utils.py
View file

@ -3080,10 +3080,9 @@ def _apply_resolved_guardrails_to_metadata(
if metadata_variable_name not in data:
data[metadata_variable_name] = {}
# Track pipeline-managed guardrails to exclude from independent execution
pipeline_managed_guardrails: set = set()
# Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode
if pipelines:
pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
data[metadata_variable_name]["_guardrail_pipelines"] = pipelines
data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails
verbose_proxy_logger.debug(
@ -3100,10 +3099,8 @@ def _apply_resolved_guardrails_to_metadata(
existing_guardrails = []
# Combine existing guardrails with policy-resolved guardrails (no duplicates)
# Exclude pipeline-managed guardrails from the flat list
combined = set(existing_guardrails)
combined.update(resolved_guardrails)
combined -= pipeline_managed_guardrails
data[metadata_variable_name]["guardrails"] = list(combined)
verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined))

262
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -5,12 +5,16 @@ Runs guardrails sequentially per pipeline step definitions, handling
pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import copy
import time
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar
from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import LOGS_GUARDRAIL_INFORMATION_MARKER
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -21,6 +25,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
independent_snapshot,
)
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -29,7 +34,14 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineStep,
PipelineStepResult,
)
from litellm.types.utils import StandardLoggingGuardrailInformation
from litellm.types.utils import GenericGuardrailAPIInputs, StandardLoggingGuardrailInformation
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
)
from litellm.proxy._types import UserAPIKeyAuth
try:
from fastapi.exceptions import HTTPException
@ -37,6 +49,121 @@ except ImportError:
HTTPException = None
class UndeliverableStreamRewrite(Exception):
def __init__(self, guardrail_name: str) -> None:
super().__init__(
f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's "
"streaming pipeline cannot deliver"
)
self.guardrail_name: Final = guardrail_name
def _tool_call_shape(tool_call: object) -> tuple[object, object]:
plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
function: Final = plain.get("function") if isinstance(plain, Mapping) else None
if not isinstance(function, Mapping):
return (None, None)
return (function.get("name"), function.get("arguments"))
def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
return None if texts is None else tuple(texts)
def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool:
return sent is not None and returned is not None and returned != sent
_GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object])
def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT:
vars(method)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the method the class body just defined
return method
class _StreamRewriteObserver(CustomGuardrail):
"""Stand-in handed to the endpoint translation in place of a streaming pipeline step's
guardrail. It records whether the guardrail returned different output than it was given,
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text
rewrites are deliverable on translations that write them back across the buffered chunks
(``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any
other translation are discarded by the executor, which releases the original chunks.
The inner guardrail's ``apply_guardrail`` already records the guardrail information
and span, so the observer's stays out of ``log_guardrail_information``."""
def __init__(self, inner: CustomGuardrail) -> None:
super().__init__(guardrail_name=inner.guardrail_name)
self.inner: Final = inner
self.rewrote_texts = False
self.rewrote_tool_calls = False
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@_logged_by_inner_guardrail
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
sent_texts: Final = _text_snapshot(inputs.get("texts"))
sent_tool_shapes: Final = _tool_call_shapes(inputs.get("tool_calls"))
outputs: Final = await self.inner.apply_guardrail(
inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
)
self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(
sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))
)
return outputs
def _prepare_hook_input(
step: PipelineStep,
callback: CustomGuardrail,
data: dict, # mutable-ok: same request-payload shape the hooks mutate
raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data
) -> tuple[dict, bool]: # mutable-ok: returns that same request-payload dict
"""Inject the step's guardrail name into metadata so should_run_guardrail() allows it,
and pick the payload the step scans: a scan_raw_request step evaluates the pristine
pre-pipeline snapshot instead of `data` (which earlier pass_data steps in this same
pipeline may have already rewritten), same reason the normal sequential/parallel
guardrail loops do this."""
if "metadata" not in data:
data["metadata"] = {} # mutable-ok: request metadata bucket, hooks mutate it
data["metadata"]["guardrails"] = [
step.guardrail
] # mutable-ok: guardrails list is part of the request-payload shape
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] # mutable-ok: request metadata shape
return hook_input, scans_raw_request
def _release_original_chunks(
guardrail_name: str,
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place
originals: Sequence[object],
) -> None:
streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives
verbose_proxy_logger.warning(
"Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming "
"pipeline cannot deliver yet; the rewrite was discarded and the original stream released",
guardrail_name,
)
class PipelineExecutor:
"""Executes guardrail pipelines with ordered, conditional step logic."""
@ -49,6 +176,8 @@ class PipelineExecutor:
call_type: str,
policy_name: str,
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step
endpoint_translation: "BaseTranslation | None" = None,
) -> PipelineExecutionResult:
"""
Execute pipeline steps sequentially with conditional actions.
@ -65,6 +194,12 @@ class PipelineExecutor:
step whose guardrail opted into ``scan_raw_request`` evaluates
the original request instead of whatever an earlier
``pass_data`` step in this same pipeline already rewrote.
streaming_chunks: buffered chunks of a completed stream. When set
(with ``endpoint_translation``), post_call steps scan the
assembled streamed output through the endpoint translation
instead of calling ``async_post_call_success_hook``.
endpoint_translation: the guardrail translation for the streamed
endpoint, resolved by the caller.
Returns:
PipelineExecutionResult with terminal action and step results
@ -89,6 +224,8 @@ class PipelineExecutor:
user_api_key_dict=user_api_key_dict,
call_type=call_type,
raw_request_snapshot=raw_request_snapshot,
streaming_chunks=streaming_chunks,
endpoint_translation=endpoint_translation,
)
duration = time.perf_counter() - start_time
@ -114,8 +251,10 @@ class PipelineExecutor:
action,
)
# Forward modified data to next step if pass_data is True
if step.pass_data and modified_data is not None:
# Forward modified data to the next step if pass_data is True;
# post_call response replacements always chain, matching the flat
# callback loop where each hook sees the previous hook's response
if modified_data is not None and (step.pass_data or mode == "post_call"):
working_data = {**working_data, **modified_data}
# Handle terminal actions
@ -129,6 +268,7 @@ class PipelineExecutor:
step_results=step_results,
error_message=error_detail,
original_exception=original_exception,
modified_data=working_data if working_data != data else None,
)
if action == "modify_response":
@ -137,6 +277,7 @@ class PipelineExecutor:
terminal_action="modify_response",
step_results=step_results,
modify_response_message=step.modify_response_message or error_detail,
modified_data=working_data if working_data != data else None,
)
# action == "next" → continue to next step
@ -144,6 +285,51 @@ class PipelineExecutor:
# Ran out of steps without a terminal action → default allow
return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
@staticmethod
async def _run_streaming_step(
step: PipelineStep,
callback: CustomGuardrail,
endpoint_translation: "BaseTranslation",
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place
hook_input: dict[str, object], # mutable-ok: same request-payload shape as data
user_api_key_dict: "UserAPIKeyAuth | None",
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Run one streaming post_call step through the endpoint translation, delivering
text rewrites on translations that support ended-stream write-back. A rewrite that
cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation
without write-back, or one the translation refused with
``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the
originals and the step passes, so the client gets the stream the merge base sent."""
observer: Final = _StreamRewriteObserver(callback)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
originals: Final = copy.deepcopy(streaming_chunks)
try:
if deliver_rewrites:
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
deliver_ended_stream_rewrites=True,
)
else:
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
)
except UndeliverableStreamRewrite:
_release_original_chunks(step.guardrail, streaming_chunks, originals)
else:
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
if not callback.records_own_guardrail_information:
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)
@staticmethod
async def _run_step(
step: PipelineStep,
@ -152,6 +338,8 @@ class PipelineExecutor:
user_api_key_dict: Any,
call_type: str,
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step
endpoint_translation: "BaseTranslation | None" = None,
) -> tuple[
Literal["pass", "fail", "error"],
dict | None,
@ -175,29 +363,13 @@ class PipelineExecutor:
verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail)
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot)
snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input))
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
if use_unified:
use_unified: Final = PipelineExecutor.supports_unified_execution(callback)
if use_unified and streaming_chunks is None:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
@ -213,6 +385,24 @@ class PipelineExecutor:
callback.mark_pre_call_hook_ran(data)
if isinstance(response, dict):
callback.mark_pre_call_hook_ran(response)
elif mode == "post_call" and streaming_chunks is not None:
if not use_unified or endpoint_translation is None:
return (
"error",
None,
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
None,
)
await PipelineExecutor._run_streaming_step(
step=step,
callback=callback,
endpoint_translation=endpoint_translation,
streaming_chunks=streaming_chunks,
hook_input=hook_input,
user_api_key_dict=user_api_key_dict,
litellm_logging_obj=data.get("litellm_logging_obj"),
)
response = None
elif mode == "post_call":
response = await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
@ -226,11 +416,19 @@ class PipelineExecutor:
# same contract as run_in_parallel/scan_raw_request elsewhere: any
# data it returned is discarded, since applying it on top of the
# raw snapshot would silently undo whatever an earlier step in
# this pipeline already did.
modified_data = None
if response is not None and isinstance(response, dict) and not scans_raw_request:
modified_data = response
return ("pass", modified_data, None, None)
# this pipeline already did. A post_call hook's non-None return is
# a replacement response (the flat callback-loop contract), carried
# under the same "response" key the step input uses.
if response is None or scans_raw_request:
return ("pass", None, None, None)
if mode == "post_call":
return (
"pass",
{"response": response},
None,
None,
) # mutable-ok: modified-data contract is a plain dict
return ("pass", response if isinstance(response, dict) else None, None, None)
except Exception as e:
if CustomGuardrail._is_guardrail_intervention(e):
@ -246,6 +444,12 @@ class PipelineExecutor:
entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:],
)
@staticmethod
def supports_unified_execution(callback: CustomGuardrail) -> bool:
"""Whether this guardrail runs through the unified apply_guardrail path,
the interface streaming pipeline execution requires."""
return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""

344
litellm/proxy/utils.py
View file

@ -11,7 +11,7 @@ import sys
import threading
import time
import traceback
from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import date, datetime, timedelta, timezone
from email.mime.multipart import MIMEMultipart
@ -139,6 +139,7 @@ from litellm.proxy.db.token_auth import (
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
resolve_endpoint_translation,
)
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
@ -449,12 +450,161 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail
)
def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]:
managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails")
return (
frozenset(cast("Collection[str]", managed)) # cast-ok: the policy engine wrote these guardrail names
if managed
else frozenset()
def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> frozenset[str]:
return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
def _pipeline_managed_guardrail_names(
data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
) -> frozenset[str]:
return _pipeline_step_guardrail_names(
tuple((policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == mode)
)
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
resolved: Final = tuple(
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
cast( # cast-ok: the resolver returns None for unknown names, filtered below
_custom_logger_compatible_callbacks_literal, callback
)
)
if isinstance(callback, str)
else callback
for callback in litellm.callbacks
)
present: Final = tuple(callback for callback in resolved if callback is not None)
guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail))
others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger
"tuple[CustomLogger, ...]",
tuple(callback for callback in present if not isinstance(callback, CustomGuardrail)),
)
return (guardrails, others)
def _merge_pipeline_metadata_bucket(
data: dict, bucket_key: str, modified_bucket_value: object
) -> None: # mutable-ok: request payload dict, written in place
if not isinstance(modified_bucket_value, dict):
return
modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed
surviving_writes: Final = {
key: value for key, value in modified_bucket.items() if key != "guardrails"
} # mutable-ok: merged into the live request metadata bucket in place
existing_bucket: Final = data.get(bucket_key)
if isinstance(existing_bucket, dict):
cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed
else:
data[bucket_key] = surviving_writes
def _merge_pipeline_metadata_writes(
data: dict, modified_data: Mapping[str, object]
) -> None: # mutable-ok: request payload dict, written in place
"""
Copy metadata-bucket writes from a pipeline's working copy back onto the request.
Post_call pipelines run step hooks against a copied request dict so the payload
already sent upstream stays untouched, but hooks record proxy-internal logging
state in the metadata buckets (``applied_guardrails`` for response headers,
``standard_logging_guardrail_information`` for spend logs), and those writes
must reach the request dict the proxy keeps reading after the pipeline returns.
The ``guardrails`` key is the executor's per-step activation flag for
``should_run_guardrail``, not a hook write, so it stays in the working copy.
"""
for bucket_key in ("metadata", "litellm_metadata"):
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
return tuple(
(policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call"
)
def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None:
if data.get("background") is not True:
return
policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data))
if not policy_names:
return
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines do not run on background responses yet; "
"the response is released ungoverned by them: %s",
", ".join(policy_names),
)
def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
unsupported: Final = tuple(
dict.fromkeys(
step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail)
)
)
if not unsupported:
return True
verbose_proxy_logger.warning(
"Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, "
"which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s",
policy_name,
", ".join(unsupported),
)
return False
def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None
def _stream_gated_guardrail_names(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> frozenset[str]:
if not _route_supports_streaming_pipelines(user_api_key_dict):
return frozenset()
return _pipeline_step_guardrail_names(
tuple(
(policy_name, pipeline)
for policy_name, pipeline in _post_call_pipelines(request_data)
if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
)
)
def _streamable_post_call_pipelines(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
"""
The post_call pipelines a streaming response can be gated through.
Streaming pipelines scan the buffered stream through the endpoint guardrail
translation of the request route, so every step's guardrail needs the
unified apply_guardrail interface and the route needs a translation. A
pipeline that cannot be run that way yet is left out and its guardrails
run on the stream on their own, the way they did before pipelines ran on
streams at all, with a warning naming the pipeline.
"""
post_call_pipelines: Final = _post_call_pipelines(request_data)
if not post_call_pipelines:
return ()
if not _route_supports_streaming_pipelines(user_api_key_dict):
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
"on their own: %s",
user_api_key_dict.request_route,
", ".join(policy_name for policy_name, _pipeline in post_call_pipelines),
)
return ()
return tuple(
(policy_name, pipeline)
for policy_name, pipeline in post_call_pipelines
if _pipeline_is_streamable(policy_name, pipeline)
)
@ -1596,7 +1746,8 @@ class ProxyLogging:
call_type: str,
event_hook: str,
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
) -> dict:
response: LLMResponseTypes | None = None,
) -> tuple[dict, LLMResponseTypes | None]: # mutable-ok: returns the request-payload dict onward
"""
Execute guardrail pipelines if any are configured for this request.
@ -1608,20 +1759,27 @@ class ProxyLogging:
``scan_raw_request`` evaluates the pristine request, not whatever an
earlier ``pass_data`` step in the same pipeline already rewrote.
Returns the (possibly modified) data dict.
Returns the (possibly modified) data dict, plus the replacement
response when a post_call pipeline step returned one (None when the
response is unchanged), matching the flat callback-loop contract.
"""
pipelines: Final = _policy_pipelines(data)
if not pipelines:
return data
return data, None
current_response = response # rebind-ok: chains each pipeline's replacement response into the next
for policy_name, pipeline in pipelines:
if pipeline.mode != event_hook:
continue
step_input: dict = (
{**data, "response": current_response} if current_response is not None else data
) # mutable-ok: same request-payload shape as data
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data=data,
data=step_input,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
policy_name=policy_name,
@ -1632,26 +1790,46 @@ class ProxyLogging:
result=result,
data=data,
policy_name=policy_name,
original_response=current_response,
)
return data
if current_response is not None and result.modified_data is not None:
current_response = result.modified_data.get("response", current_response)
return data, current_response if current_response is not response else None
@staticmethod
def _handle_pipeline_result(
result: PipelineExecutionResult,
data: dict,
policy_name: str,
original_response: "LLMResponseTypes | Sequence[object] | None" = None,
) -> dict:
"""
Handle a PipelineExecutionResult — allow, block, or modify_response.
Returns data dict if allowed, raises on block/modify_response.
``original_response`` is set on the post_call path, where the request
payload (already sent upstream) must stay untouched; a replacement
response carried in ``modified_data`` is adopted by the caller, and
metadata-bucket writes (applied guardrails, guardrail logging info)
are merged back so headers and spend logs still see them, on block
and modify_response too, so failure spend records keep guardrail
cost and status. On the
streaming path it is the buffered chunk list, carried into
``ModifyResponseException.original_response`` for usage reporting.
"""
if result.terminal_action == "allow":
if result.modified_data is not None:
data.update(result.modified_data)
if original_response is None:
data.update(result.modified_data)
else:
_merge_pipeline_metadata_writes(data, result.modified_data)
return data
if result.modified_data is not None:
_merge_pipeline_metadata_writes(data, result.modified_data)
if result.terminal_action == "block":
original_exception: Final = result.original_exception
if original_exception is not None and not _exception_changes_request_flow(original_exception):
@ -1689,6 +1867,7 @@ class ProxyLogging:
request_data=data,
guardrail_name=f"pipeline:{policy_name}",
detection_info=None,
original_response=original_response,
)
return data
@ -1805,8 +1984,10 @@ class ProxyLogging:
)
try:
_warn_background_skips_post_call_pipelines(data)
# Execute guardrail pipelines before the normal callback loop
data = await self._maybe_execute_pipelines(
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
@ -1815,7 +1996,7 @@ class ProxyLogging:
)
# Get pipeline-managed guardrails to skip in normal loop
pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
caps: Final = ProxyLogging._callback_capabilities()
# Skip the per-request callback walk entirely when nothing in
@ -2793,36 +2974,35 @@ class ProxyLogging:
from litellm.proxy.proxy_server import llm_router
from litellm.types.guardrails import GuardrailEventHooks
guardrail_callbacks: Final[list[CustomGuardrail]] = []
other_callbacks: Final[list[CustomLogger]] = []
_, pipeline_response = await self._maybe_execute_pipelines(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
event_hook="post_call",
response=response,
)
if pipeline_response is not None:
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
try:
for callback in litellm.callbacks:
_callback: CustomLogger | None = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
cast(_custom_logger_compatible_callbacks_literal, callback)
)
else:
_callback = callback
if _callback is not None:
if isinstance(_callback, CustomGuardrail):
guardrail_callbacks.append(_callback)
else:
other_callbacks.append(_callback)
############## Handle Guardrails ########################################
#############################################################################
# Merge model-level guardrails before checking which guardrails to run
guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False)
callback
for callback in guardrail_callbacks
if getattr(callback, "run_in_parallel", False)
and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed)
)
for callback in guardrail_callbacks:
# Main - V2 Guardrails implementation
if callback.guardrail_name and callback.guardrail_name in pipeline_managed:
continue
if getattr(callback, "run_in_parallel", False):
continue
@ -3119,11 +3299,16 @@ class ProxyLogging:
# dict lookups + llm_router.get_deployment() per callback per chunk.
_cached_guardrail_data: dict | None = None
_guardrail_data_computed = False
pipeline_gated: Final = (
_stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
)
for callback in litellm.callbacks:
try:
_callback: CustomLogger | None = None
if isinstance(callback, CustomGuardrail):
if callback.guardrail_name in pipeline_gated:
continue
# Main - V2 Guardrails implementation
from litellm.types.guardrails import GuardrailEventHooks
@ -3180,12 +3365,13 @@ class ProxyLogging:
1. /chat/completions
"""
caps: Final = ProxyLogging._callback_capabilities()
post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
# Fast path: no real overrides. Internal proxy CustomLogger callbacks
# (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default
# ``async for chunk: yield chunk`` body, so wrapping the iterator
# through each of them adds N pass-through trampolines per chunk for
# zero behavior change. Skip the chain entirely and stream through.
if not caps.iterator_overrides:
if not caps.iterator_overrides and not post_call_pipelines:
try:
async for chunk in response:
yield chunk
@ -3205,8 +3391,11 @@ class ProxyLogging:
current_response = response
stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines)
for resolved_callback, kind in caps.iterator_overrides:
if isinstance(resolved_callback, CustomGuardrail):
if resolved_callback.guardrail_name in pipeline_gated_names:
continue
if (
resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call)
is not True
@ -3246,6 +3435,14 @@ class ProxyLogging:
),
)
if post_call_pipelines:
current_response = self._pipeline_gated_stream(
response=current_response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
pipelines=post_call_pipelines,
)
try:
async for chunk in current_response:
yield chunk
@ -3261,6 +3458,81 @@ class ProxyLogging:
# we reach this point the metadata is fully populated.
ProxyLogging._fire_deferred_stream_logging(request_data)
async def _pipeline_gated_stream(
self,
response: "AsyncGenerator[object, None]",
user_api_key_dict: UserAPIKeyAuth,
request_data: dict, # mutable-ok: same request-payload shape the hooks mutate
pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
) -> "AsyncGenerator[Any, None]":
"""
Execute post_call policy pipelines against a streamed response.
Buffers the whole stream (nothing reaches the client until every
pipeline allows it), then runs each pipeline's steps against the
assembled output through the endpoint guardrail translation, the same
machinery flat post_call guardrails use at end of stream. An allow
releases the buffered chunks: verbatim when no guardrail rewrote the
output, rewritten in place when one rewrote text and the translation
delivers ended-stream rewrites (later steps then re-scan the rewritten
chunks, so rewrites chain). A rewrite the translation cannot deliver
yet (a tool-call rewrite, or a text rewrite on a route without
write-back) is discarded by the executor and the original chunks are
released, as is a buffered shape no translation resolves; a block or
modify_response terminates with the translation's block chunks or the
raised error.
"""
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
async for item in response:
buffered.append(item)
if not buffered:
return
resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
if resolved is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; "
"the stream is released ungoverned by them: %s",
", ".join(policy_name for policy_name, _pipeline in pipelines),
)
for buffered_item in buffered:
yield buffered_item
return
call_type, endpoint_translation = resolved
for policy_name, pipeline in pipelines:
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode="post_call",
data=request_data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
policy_name=policy_name,
streaming_chunks=buffered,
endpoint_translation=endpoint_translation,
)
try:
ProxyLogging._handle_pipeline_result(
result, data=request_data, policy_name=policy_name, original_response=buffered
)
except ModifyResponseException as e:
if e.original_response is None:
e.original_response = buffered
async for block_chunk in unified_guardrail.handle_streaming_block(
e, endpoint_translation, stream_started=False, responses_so_far=()
):
yield block_chunk
return
except HTTPException as e:
async for error_chunk in unified_guardrail.emit_streaming_http_error(
e, call_type, buffered, request_data
):
yield error_chunk
return
for buffered_item in buffered:
yield buffered_item
@staticmethod
def _fire_deferred_stream_logging(request_data: dict) -> None:
"""

111
tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -264,6 +264,117 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
# Should return the responses unchanged
assert result == responses_so_far
@staticmethod
def _ended_sse_chunks() -> list:
events = [
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@staticmethod
def _masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs.get("texts", [])]}
return MaskWorld(guardrail_name="test")
@staticmethod
def _delta_texts(chunks: list) -> list:
texts = []
for chunk in chunks:
for line in chunk.decode().split("\n"):
if not line.startswith("data:"):
continue
data = json.loads(line[len("data:") :].strip())
if data.get("type") == "content_block_delta":
texts.append(data["delta"]["text"])
return texts
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert self._delta_texts(chunks) == ["hello [MASKED]", ""]
raw = b"".join(chunks).decode()
assert "event: message_start" in raw and "event: message_stop" in raw
assert '"stop_reason": "end_turn"' in raw
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()
original = [bytes(chunk) for chunk in chunks]
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
original = [bytes(chunk) for chunk in chunks]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_rewrite_without_delivery_expected_does_not_raise(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
original = [bytes(chunk) for chunk in chunks]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert result is chunks
assert chunks == original
class TestAnthropicMessagesHandlerInputProcessing:
"""Test input processing preserves litellm_metadata for dynamic guardrails."""

148
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -1074,6 +1074,154 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
# Should return the responses
assert result == responses_so_far
@staticmethod
def _ended_stream_chunks() -> list:
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
return [
ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=0, delta=Delta(content="Hello"), finish_reason=None)],
),
ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")],
),
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_text_back_into_chunks(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_stream_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks[0].choices[0].delta.content == "HELLO WORLD"
assert chunks[1].choices[0].delta.content in (None, "")
assert chunks[1].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert chunks[0].choices[0].delta.content == "Hello"
assert chunks[1].choices[0].delta.content == " world"
assert chunks[1].choices[0].finish_reason == "stop"
@staticmethod
def _two_choice_stream_chunks() -> list:
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)],
)
return [
chunk(0, "safe "),
chunk(1, "hello "),
chunk(0, "text", "stop"),
chunk(1, "world", "stop"),
]
@staticmethod
def _world_masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
texts = inputs.get("texts", [])
return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
return MaskWorld(guardrail_name="test-mask")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_stream_chunks()
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._world_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_stream_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)],
)
chunks = [chunk("hello ", None), chunk("world", "stop")]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._world_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks[0].choices[0].delta.content == "hello [MASKED]"
assert chunks[1].choices[0].delta.content in (None, "")
class TestUndecoratedGuardrailIsRecorded:
"""LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail

203
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -1128,6 +1128,209 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
output_text = result[-1]["response"]["output"][0]["content"][0]["text"]
assert output_text == original_text
@staticmethod
def _ended_stream_events() -> List[dict]:
content = [{"type": "output_text", "text": "hello world"}]
item = {
"type": "message",
"id": "msg_123",
"status": "completed",
"role": "assistant",
"content": content,
}
return [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
{
"type": "response.content_part.done",
"output_index": 0,
"content_index": 0,
"part": {"type": "output_text", "text": "hello world"},
},
{"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"model": "gpt-4o",
"output": [{**item, "content": [dict(c) for c in content]}],
"status": "completed",
},
},
]
@staticmethod
def _masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
texts = inputs.get("texts", [])
return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
return MaskWorld(guardrail_name="test-mask")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_all_stream_events(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["delta"] == ""
assert events[2]["text"] == "hello [MASKED]"
assert events[3]["part"]["text"] == "hello [MASKED]"
assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
@pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
events[-1]["type"] = terminal_type
events[-1]["response"]["status"] = terminal_type.split(".")[-1]
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["delta"] == ""
assert events[2]["text"] == "hello [MASKED]"
assert events[3]["part"]["text"] == "hello [MASKED]"
assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_fallback_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"},
]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_output_item_done_last_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_output_item_done_last_scans_text_with_delivery_expected(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
guardrail = MockRecordingGuardrail(guardrail_name="test")
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]]
@pytest.mark.asyncio
async def test_output_item_done_last_without_delivery_expected_skips_text(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
guardrail = MockRecordingGuardrail(guardrail_name="test")
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert result is events
assert guardrail.seen_inputs == []
@pytest.mark.asyncio
async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self):
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
]
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
)
assert result is events
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
)
assert events[0]["delta"] == "hello "
assert events[1]["delta"] == "world"
assert events[2]["text"] == "hello world"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_failed_stream_scans_delta_text(self):
"""A stream ending in response.failed has text only in delta events; the

2
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -1034,7 +1034,7 @@ class TestStreamingTransform:
)
emitted = []
async for item in handler._emit_streaming_http_error(
async for item in handler.emit_streaming_http_error(
exc,
call_type=CallTypes.asend_message.value,
responses_so_far=[{"id": "req-1"}],

244
tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
View file

@ -5,6 +5,7 @@ Uses mock guardrails to validate pipeline execution without external services.
"""
import copy
import logging
from typing import Literal
from unittest.mock import MagicMock
@ -17,7 +18,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeGuardrail,
)
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
@ -1052,3 +1053,244 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch):
assert outcome == "pass"
assert guardrail.native_pre_call_ran is True
assert "guardrail_to_apply" not in data
class _TextReturningGuardrail(CustomGuardrail):
def __init__(self, returned_texts):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
self.returned_texts = returned_texts
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": self.returned_texts}
class _TextTranslation:
delivers_ended_stream_text_rewrites = False
def __init__(self):
self.seen_guardrail_names = []
async def process_output_streaming_response(
self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
):
self.seen_guardrail_names.append(guardrail_to_apply.guardrail_name)
await guardrail_to_apply.apply_guardrail(
inputs={"texts": ["hello world"]},
request_data=request_data or {},
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
class _WritingTranslation:
"""Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the
chat/Responses/Messages handlers do on an ended stream."""
delivers_ended_stream_text_rewrites = True
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj=None,
user_api_key_dict=None,
request_data=None,
deliver_ended_stream_rewrites=False,
):
assert deliver_ended_stream_rewrites is True
outputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]},
request_data=request_data or {},
input_type="response",
logging_obj=litellm_logging_obj,
)
responses_so_far[0]["text"] = outputs["texts"][0]
responses_so_far[0]["tool_call"] = outputs["tool_calls"][0]
return responses_so_far
class _RefusingTranslation:
delivers_ended_stream_text_rewrites = True
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj=None,
user_api_key_dict=None,
request_data=None,
deliver_ended_stream_rewrites=False,
):
responses_so_far[0]["text"] = "half-written"
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name)
def _chunk():
return {"text": "hello world", "tool_call": {"function": {"name": "lookup", "arguments": '{"ssn": "123"}'}}}
async def _run_streaming_step(translation, streaming_chunks=None):
chunks = [object()] if streaming_chunks is None else streaming_chunks
return await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")],
mode="post_call",
data={"model": "m"},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="p",
streaming_chunks=chunks,
endpoint_translation=translation,
)
def _assert_passed_with_discard_warning(result, caplog):
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
translation = _TextTranslation()
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(translation, chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
assert translation.seen_guardrail_names == ["masker"]
@pytest.mark.asyncio
async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))])
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_TextTranslation())
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert not any("discarded" in record.getMessage() for record in caplog.records)
class _InPlaceMutatingGuardrail(CustomGuardrail):
"""Rewrites like bedrock/presidio do: rebinds inputs["texts"] on the dict it was handed
and returns that same dict, so a post-call comparison against inputs sees no change."""
def __init__(self):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
inputs["texts"] = ["hello [MASKED]"]
return inputs
@pytest.mark.asyncio
async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_InPlaceMutatingGuardrail()])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_TextTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
class _TextAndToolCallRewritingGuardrail(CustomGuardrail):
def __init__(self, rewrite_tool_call):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
self.rewrite_tool_call = rewrite_tool_call
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
tool_calls = (
[{"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}]
if self.rewrite_tool_call
else inputs["tool_calls"]
)
return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": tool_calls}
@pytest.mark.asyncio
async def test_streaming_step_delivers_text_rewrite_through_writing_translation(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=False)])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_WritingTranslation(), chunks)
assert result.terminal_action == "allow"
assert chunks[0]["text"] == "hello [MASKED]"
assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "123"}'
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_WritingTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
class _BlockingStreamGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
raise HTTPException(status_code=400, detail={"error": "output blocked"})
def _recorded_guardrail_statuses(result):
return [
entry["guardrail_status"]
for entry in result.modified_data["metadata"]["standard_logging_guardrail_information"]
]
@pytest.mark.asyncio
async def test_streaming_step_records_guardrail_information_once_on_mask(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
assert result.terminal_action == "allow"
assert _recorded_guardrail_statuses(result) == ["success"]
@pytest.mark.asyncio
async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_header(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
@pytest.mark.asyncio
async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()])
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
assert [step.outcome for step in result.step_results] == ["fail"]
assert _recorded_guardrail_statuses(result) == ["guardrail_intervened"]
@pytest.mark.asyncio
async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewrite(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_RefusingTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]

42
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -4148,6 +4148,48 @@ async def test_add_guardrails_from_policy_engine():
attachment_registry._initialized = False
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps():
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import (
GuardrailPipeline,
PipelineStep,
Policy,
PolicyAttachment,
PolicyGuardrails,
)
data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "metadata": {}}
policy_registry = get_policy_registry()
policy_registry._policies = {
"response-governance": Policy(
guardrails=PolicyGuardrails(add=["pii_blocker"]),
pipeline=GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii_blocker")]),
),
}
policy_registry._initialized = True
attachment_registry = get_attachment_registry()
attachment_registry._attachments = [PolicyAttachment(policy="response-governance", scope="*")]
attachment_registry._initialized = True
try:
await add_guardrails_from_policy_engine(
data=data,
metadata_variable_name="metadata",
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
finally:
policy_registry._policies = {}
policy_registry._initialized = False
attachment_registry._attachments = []
attachment_registry._initialized = False
assert data["metadata"]["guardrails"] == ["pii_blocker"]
assert data["metadata"]["_pipeline_managed_guardrails"] == {"pii_blocker"}
assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"]
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data():
"""

1254
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

File diff suppressed because it is too large Load diff

2
tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -681,7 +681,7 @@ async def test_scan_raw_request_snapshot_taken_before_pipelines(
for msg in data.get("messages", []):
if "SECRET" in msg.get("content", ""):
msg["content"] = msg["content"].replace("SECRET", "[REDACTED]")
return data
return data, None
monkeypatch.setattr(ProxyLogging, "_maybe_execute_pipelines", fake_pipelines)
monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(scan_raw_request=True)])