mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(policy_engine): apply post_call pipeline text rewrites on streams
Buffered streams governed by post_call policy pipelines now deliver text rewrites back into the stream per surface (chat SSE, responses SSE, anthropic messages SSE) instead of rejecting the request with a 400 upfront. Rewrites chain across pipeline steps; tool-call rewrites and translations without stream write-back still withhold the stream.
This commit is contained in:
parent
99286cd254
commit
673d1743a6
11 changed files with 711 additions and 130 deletions
|
|
@ -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, cast
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
|
@ -120,6 +121,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()
|
||||
|
|
@ -931,11 +934,14 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict | None = None,
|
||||
deliver_ended_stream_rewrites: bool = False,
|
||||
) -> list[Any]:
|
||||
"""
|
||||
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).
|
||||
"""
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
|
||||
|
|
@ -982,6 +988,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
|
||||
|
|
@ -1093,6 +1108,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_string_so_far(self, responses_so_far: list[Any]) -> str:
|
||||
"""
|
||||
Parse streaming responses and extract accumulated text content.
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from abc import ABC, abstractmethod
|
||||
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 litellm.integrations.custom_guardrail import (
|
||||
|
|
@ -33,6 +33,13 @@ class StreamTransformSink:
|
|||
|
||||
|
||||
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 instead of withholding the
|
||||
stream. Tool-call rewrites stay undeliverable everywhere."""
|
||||
|
||||
@staticmethod
|
||||
def transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict: Any | None,
|
||||
|
|
@ -113,6 +120,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.
|
||||
|
|
@ -120,6 +128,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
|
||||
|
||||
|
|
|
|||
|
|
@ -61,6 +61,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.
|
||||
|
|
@ -440,6 +442,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
user_api_key_dict: Any | 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.
|
||||
|
|
@ -454,6 +457,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.
|
||||
|
|
@ -479,6 +486,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(
|
||||
|
|
@ -489,10 +497,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
user_api_key_dict: Any | 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."""
|
||||
# check if the stream has ended
|
||||
has_stream_ended = False
|
||||
for chunk in responses_so_far:
|
||||
|
|
@ -501,20 +511,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
break
|
||||
|
||||
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
|
||||
|
|
@ -591,6 +595,38 @@ 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: object,
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _accumulate_string_content_by_choice_index(
|
||||
responses_so_far: list["ModelResponseStream"],
|
||||
|
|
@ -922,6 +958,41 @@ 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, ...],
|
||||
) -> None:
|
||||
"""Write ended-stream guardrail text rewrites back across the buffered
|
||||
chunks: each rewritten choice's full text lands in its 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."""
|
||||
post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response)
|
||||
changed: Final = tuple(
|
||||
(choice_idx, after)
|
||||
for choice_idx, (before, after) in enumerate(zip(pre_guardrail_texts, post_guardrail_texts))
|
||||
if before is not None and after is not None and after != before
|
||||
)
|
||||
if not changed:
|
||||
return
|
||||
await self._apply_guardrail_responses_to_output_streaming(
|
||||
responses=responses_so_far,
|
||||
guardrailed_texts=[
|
||||
after for _choice_idx, after in changed
|
||||
], # mutable-ok: the callee's signature predates this change and takes lists
|
||||
task_mappings=[
|
||||
(choice_idx, None) for choice_idx, _after in changed
|
||||
], # mutable-ok: the callee's signature predates this change and takes lists
|
||||
)
|
||||
|
||||
async def _apply_guardrail_responses_to_output_streaming(
|
||||
self,
|
||||
responses: list["ModelResponseStream"],
|
||||
|
|
|
|||
|
|
@ -28,7 +28,9 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
|
|||
- text: str
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import chain, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
|
||||
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
|
||||
|
|
@ -91,6 +93,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.
|
||||
|
|
@ -482,6 +486,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict | None = None,
|
||||
deliver_ended_stream_rewrites: bool = False,
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Process output streaming response by applying guardrails to text content.
|
||||
|
|
@ -493,7 +498,11 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
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 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 completed response too, so a client reading deltas sees
|
||||
the rewrite instead of the raw model output.
|
||||
"""
|
||||
if not responses_so_far:
|
||||
return responses_so_far
|
||||
|
|
@ -562,6 +571,19 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
responses=guardrailed_texts,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
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
|
||||
|
||||
|
|
@ -607,6 +629,63 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
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[ResponsesStreamChunk]) -> bool:
|
||||
"""
|
||||
Check if the streaming has ended.
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ if TYPE_CHECKING:
|
|||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
try:
|
||||
from fastapi.exceptions import HTTPException
|
||||
|
|
@ -44,7 +45,8 @@ except ImportError:
|
|||
class UndeliverableStreamRewrite(Exception):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(
|
||||
f"Guardrail '{guardrail_name}' rewrote the streamed response, which streaming pipelines cannot deliver"
|
||||
f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's "
|
||||
"streaming pipeline cannot deliver"
|
||||
)
|
||||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
|
@ -57,28 +59,31 @@ def _tool_call_shape(tool_call: object) -> tuple[object, object]:
|
|||
return (function.get("name"), function.get("arguments"))
|
||||
|
||||
|
||||
def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool:
|
||||
return sent is not None and returned is not None and list(returned) != list(sent)
|
||||
def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
|
||||
return None if texts is None else tuple(texts)
|
||||
|
||||
|
||||
def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool:
|
||||
if sent is None or returned is None:
|
||||
return False
|
||||
return [_tool_call_shape(tool_call) for tool_call in returned] != [
|
||||
_tool_call_shape(tool_call) for tool_call in sent
|
||||
]
|
||||
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
|
||||
|
||||
|
||||
class _StreamRewriteObserver(CustomGuardrail):
|
||||
"""Stand-in handed to the endpoint translation in place of a streaming pipeline step's
|
||||
guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate
|
||||
withholds the stream whenever the guardrail returned different output than it was given,
|
||||
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime."""
|
||||
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 make the gate withhold the stream."""
|
||||
|
||||
def __init__(self, inner: CustomGuardrail) -> None:
|
||||
super().__init__(guardrail_name=inner.guardrail_name)
|
||||
self.inner: Final = inner
|
||||
self.rewrote = False
|
||||
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()
|
||||
|
|
@ -90,13 +95,14 @@ class _StreamRewriteObserver(CustomGuardrail):
|
|||
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 = (
|
||||
self.rewrote
|
||||
or _rewrote_texts(inputs.get("texts"), outputs.get("texts"))
|
||||
or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls"))
|
||||
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
|
||||
|
||||
|
|
@ -250,6 +256,41 @@ class PipelineExecutor:
|
|||
modified_data=working_data if working_data != data else None,
|
||||
)
|
||||
|
||||
@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 and raising
|
||||
``UndeliverableStreamRewrite`` for any rewrite that cannot reach the client."""
|
||||
observer: Final = _StreamRewriteObserver(callback)
|
||||
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
|
||||
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,
|
||||
)
|
||||
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
|
||||
raise UndeliverableStreamRewrite(step.guardrail)
|
||||
|
||||
@staticmethod
|
||||
async def _run_step(
|
||||
step: PipelineStep,
|
||||
|
|
@ -312,16 +353,15 @@ class PipelineExecutor:
|
|||
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
|
||||
None,
|
||||
)
|
||||
observer: Final = _StreamRewriteObserver(callback)
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=streaming_chunks,
|
||||
guardrail_to_apply=observer,
|
||||
litellm_logging_obj=data.get("litellm_logging_obj"),
|
||||
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,
|
||||
request_data=hook_input,
|
||||
litellm_logging_obj=data.get("litellm_logging_obj"),
|
||||
)
|
||||
if observer.rewrote:
|
||||
raise UndeliverableStreamRewrite(step.guardrail)
|
||||
response = None
|
||||
elif mode == "post_call":
|
||||
response = await target.async_post_call_success_hook(
|
||||
|
|
|
|||
|
|
@ -492,14 +492,6 @@ def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
|
|||
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
|
||||
|
||||
|
||||
def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool:
|
||||
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
|
||||
if callback is None:
|
||||
return False
|
||||
transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only")
|
||||
return callback.rewrites_streamed_output() or transform_mode == "incremental_diff"
|
||||
|
||||
|
||||
class _PipelineErrorBody(TypedDict):
|
||||
message: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
|
@ -516,9 +508,10 @@ def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) -
|
|||
"error": {
|
||||
"message": (
|
||||
f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail "
|
||||
f"'{guardrail_name}' rewrote the streamed output, and streaming pipelines cannot deliver "
|
||||
"rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add "
|
||||
"applies it to streamed output."
|
||||
f"'{guardrail_name}' rewrote the streamed output in a way this endpoint's streaming "
|
||||
"pipeline cannot deliver (a tool-call rewrite, or a text rewrite on a route without "
|
||||
"stream write-back). Retry with stream=false, or drop it from the pipeline steps so "
|
||||
"guardrails.add applies it to streamed output."
|
||||
),
|
||||
"type": "guardrail_pipeline_error",
|
||||
"policies": (policy_name,),
|
||||
|
|
@ -535,13 +528,13 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap
|
|||
Background responses skip the post_call hooks entirely, so a pipeline
|
||||
governing one would silently never execute. Streaming responses execute
|
||||
pipelines against the buffered stream through the endpoint guardrail
|
||||
translation of the request route, releasing the buffered chunks on allow.
|
||||
That needs every step's guardrail to support the unified apply_guardrail
|
||||
interface and to only allow or block (a step that rewrites streamed
|
||||
content, via mask_response_content, a MASK action, or
|
||||
streaming_transform_mode=incremental_diff, would have its rewrite silently
|
||||
dropped), and needs the route to have a translation at all; anything else
|
||||
keeps the 400 rather than letting ungoverned output stream through.
|
||||
translation of the request route, releasing the buffered chunks on allow
|
||||
(rewritten in place when a guardrail rewrote text and the translation
|
||||
delivers ended-stream rewrites; a rewrite the translation cannot deliver
|
||||
fails closed at runtime instead). That needs every step's guardrail to
|
||||
support the unified apply_guardrail interface, and needs the route to have
|
||||
a translation at all; anything else keeps the 400 rather than letting
|
||||
ungoverned output stream through.
|
||||
"""
|
||||
is_stream: Final = data.get("stream") is True
|
||||
is_background: Final = data.get("background") is True
|
||||
|
|
@ -586,25 +579,6 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap
|
|||
}
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=unsupported_detail)
|
||||
rewriting_guardrails: Final = tuple(
|
||||
guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail)
|
||||
)
|
||||
if rewriting_guardrails:
|
||||
rewriting_detail: Final[_PipelineErrorDetail] = {
|
||||
"error": {
|
||||
"message": (
|
||||
"Policies with post_call guardrail pipelines cannot govern streaming responses "
|
||||
"because these pipeline guardrails rewrite streamed content (mask_response_content, "
|
||||
"a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release "
|
||||
f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop "
|
||||
"them from the pipeline steps so guardrails.add applies them to streamed output."
|
||||
),
|
||||
"type": "guardrail_pipeline_error",
|
||||
"policies": post_call_policies,
|
||||
"guardrails": rewriting_guardrails,
|
||||
}
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=rewriting_detail)
|
||||
route: Final = user_api_key_dict.request_route
|
||||
if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None:
|
||||
return
|
||||
|
|
@ -3485,12 +3459,13 @@ class ProxyLogging:
|
|||
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; a step whose guardrail rewrote
|
||||
the output withholds the stream with a 400 instead, since no
|
||||
translation rewrites every buffered chunk consistently and some
|
||||
rewrites (Bedrock's ANONYMIZED action, for one) are only decided at
|
||||
runtime; a block or modify_response terminates with the translation's
|
||||
block chunks or the raised error.
|
||||
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
|
||||
(a tool-call rewrite, or a text rewrite on a route without write-back)
|
||||
withholds the stream with a 400; 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:
|
||||
|
|
|
|||
|
|
@ -263,6 +263,71 @@ 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
|
||||
|
||||
|
||||
class TestAnthropicMessagesHandlerInputProcessing:
|
||||
"""Test input processing preserves litellm_metadata for dynamic guardrails."""
|
||||
|
|
|
|||
|
|
@ -1073,6 +1073,61 @@ 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"
|
||||
|
||||
|
||||
class TestGetStructuredMessages:
|
||||
"""Test the get_structured_messages method."""
|
||||
|
|
|
|||
|
|
@ -1104,6 +1104,89 @@ 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
|
||||
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]"
|
||||
|
||||
|
||||
class TestGetStructuredMessages:
|
||||
"""Test the get_structured_messages method for Responses API handler."""
|
||||
|
|
|
|||
|
|
@ -823,6 +823,8 @@ class _TextReturningGuardrail(CustomGuardrail):
|
|||
|
||||
|
||||
class _TextTranslation:
|
||||
delivers_ended_stream_text_rewrites = False
|
||||
|
||||
def __init__(self):
|
||||
self.seen_guardrail_names = []
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
ModifyResponseException,
|
||||
)
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
|
||||
from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
|
||||
|
|
@ -1441,7 +1441,7 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni
|
|||
("guardrail_config", {"streaming_transform_mode": "incremental_diff"}),
|
||||
],
|
||||
)
|
||||
async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content(
|
||||
async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_rewrites_streamed_content(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
|
|
@ -1450,24 +1450,21 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_
|
|||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
guardrails_only=True,
|
||||
)
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
guardrails_only=True,
|
||||
)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
|
||||
assert "rewrite streamed content" in info.value.detail["error"]["message"]
|
||||
assert seen.get("count") is None
|
||||
assert out is not None
|
||||
assert out.get("stream") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)])
|
||||
async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected
|
||||
@pytest.mark.parametrize("action", [ContentFilterAction.MASK, ContentFilterAction.BLOCK])
|
||||
async def test_pre_call_hook_allows_streaming_when_content_filter_step_masks_or_blocks(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, action
|
||||
):
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="gr-post",
|
||||
|
|
@ -1478,25 +1475,15 @@ async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_mas
|
|||
data = _post_call_pipeline_data(stream=True)
|
||||
user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
|
||||
if not rejected:
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
|
||||
)
|
||||
assert out is not None and out.get("stream") is True
|
||||
return
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
|
||||
)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
|
||||
assert "a MASK action" in info.value.detail["error"]["message"]
|
||||
assert out is not None and out.get("stream") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks(
|
||||
async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
guardrail = ContentFilterGuardrail(
|
||||
|
|
@ -1508,13 +1495,11 @@ async def test_pre_call_hook_rejects_streaming_when_content_filter_category_mask
|
|||
data = _post_call_pipeline_data(stream=True)
|
||||
user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
|
||||
)
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True
|
||||
)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
|
||||
assert out is not None and out.get("stream") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1618,17 +1603,10 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")])
|
||||
@pytest.mark.parametrize(
|
||||
"make_chunks, transform",
|
||||
[
|
||||
(_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}),
|
||||
(_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}),
|
||||
],
|
||||
ids=["texts", "tool_calls"],
|
||||
)
|
||||
async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error
|
||||
async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewrite(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error
|
||||
):
|
||||
transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731
|
||||
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error)
|
||||
|
|
@ -1638,7 +1616,7 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
|
|||
async def _drain() -> None:
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
|
||||
response=_async_chunk_iter(make_chunks()),
|
||||
response=_async_chunk_iter(_tool_call_stream_chunks()),
|
||||
request_data=data,
|
||||
):
|
||||
delivered.append(item)
|
||||
|
|
@ -1655,6 +1633,81 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite(
|
|||
assert "stream=false" in error["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_pipeline_delivers_runtime_text_rewrite(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
|
||||
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = _stream_chunks()
|
||||
|
||||
delivered = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
|
||||
response=_async_chunk_iter(chunks),
|
||||
request_data=data,
|
||||
)
|
||||
]
|
||||
|
||||
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
|
||||
assert delivered[0].choices[0].delta.content == "hello [MASKED]"
|
||||
assert delivered[1].choices[0].delta.content in (None, "")
|
||||
assert delivered[1].choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_pipeline_chains_text_rewrites_across_steps(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
second_step_saw: Dict[str, Any] = {}
|
||||
|
||||
class FirstMask(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["texts"]]}
|
||||
|
||||
class SecondMask(CustomGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
second_step_saw["texts"] = list(inputs["texts"])
|
||||
return {**inputs, "texts": [text.replace("hello", "[GREETING]") for text in inputs["texts"]]}
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[
|
||||
FirstMask(guardrail_name="gr-first", event_hook=GuardrailEventHooks.post_call, default_on=False),
|
||||
SecondMask(guardrail_name="gr-second", event_hook=GuardrailEventHooks.post_call, default_on=False),
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
pipeline = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[
|
||||
PipelineStep(guardrail="gr-first", on_pass="next", on_fail="block"),
|
||||
PipelineStep(guardrail="gr-second", on_pass="allow", on_fail="block"),
|
||||
],
|
||||
)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)]
|
||||
chunks = _stream_chunks()
|
||||
|
||||
delivered = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
|
||||
response=_async_chunk_iter(chunks),
|
||||
request_data=data,
|
||||
)
|
||||
]
|
||||
|
||||
assert second_step_saw["texts"] == ["hello [MASKED]"]
|
||||
assert delivered[0].choices[0].delta.content == "[GREETING] [MASKED]"
|
||||
assert delivered[1].choices[0].delta.content in (None, "")
|
||||
assert delivered[1].choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"make_chunks, transform",
|
||||
|
|
@ -1761,6 +1814,79 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated
|
|||
assert not any(item is chunk for item in delivered for chunk in chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
|
||||
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = _anthropic_sse_chunks()
|
||||
|
||||
delivered = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
|
||||
response=_async_chunk_iter(chunks),
|
||||
request_data=data,
|
||||
)
|
||||
]
|
||||
|
||||
raw = b"".join(delivered).decode()
|
||||
assert "hello [MASKED]" in raw
|
||||
assert "hello world" not in raw
|
||||
assert raw.count("event: content_block_delta") == 1
|
||||
for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"):
|
||||
assert f"event: {expected_event}" in raw
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_write_back(monkeypatch):
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
|
||||
|
||||
class NoWriteBackTranslation(BaseTranslation):
|
||||
async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj):
|
||||
return data
|
||||
|
||||
async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj, **kwargs):
|
||||
return response
|
||||
|
||||
async def process_output_streaming_response(
|
||||
self,
|
||||
responses_so_far,
|
||||
guardrail_to_apply,
|
||||
litellm_logging_obj,
|
||||
user_api_key_dict=None,
|
||||
request_data=None,
|
||||
stream_transform_sink=None,
|
||||
deliver_ended_stream_rewrites=False,
|
||||
):
|
||||
assert deliver_ended_stream_rewrites is False
|
||||
await guardrail_to_apply.apply_guardrail(
|
||||
inputs={"texts": ["hello world"]},
|
||||
request_data=request_data or {},
|
||||
input_type="response",
|
||||
)
|
||||
return responses_so_far
|
||||
|
||||
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
|
||||
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
|
||||
|
||||
with pytest.raises(UndeliverableStreamRewrite):
|
||||
await PipelineExecutor.execute_steps(
|
||||
steps=[PipelineStep(guardrail="gr-post", on_pass="allow", on_fail="block")],
|
||||
mode="post_call",
|
||||
data={"metadata": {}},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
call_type="acompletion",
|
||||
policy_name="response-governance",
|
||||
streaming_chunks=_stream_chunks(),
|
||||
endpoint_translation=NoWriteBackTranslation(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_pipeline_gates_without_iterator_overrides(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue