mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(policy_engine): fail open on streaming shapes post_call pipelines cannot govern yet
A post_call pipeline now releases the original stream instead of refusing the request on every shape it has no handler for: a background request, a pipeline guardrail without the unified apply_guardrail interface, a route with no endpoint translation, a buffered stream no translation resolves, and a rewrite the translation cannot write back (tool-call edits, text edits on translations without write-back, n>1 chat, an unended Anthropic stream, a Responses dump with no event envelope). Each case logs a warning naming the policy and guardrail. Real blocks and writable text masks are unchanged.
This commit is contained in:
mateo-berri
parent
0d5ea553da
commit
192ea9ec80
8 changed files with 437 additions and 294 deletions
|
|
@ -1027,7 +1027,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
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 fails closed instead.
|
||||
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
|
||||
|
||||
|
|
|
|||
|
|
@ -56,8 +56,9 @@ class BaseTranslation(ABC):
|
|||
"""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."""
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -1015,7 +1015,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
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 fails closed."""
|
||||
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
|
||||
|
|
@ -1030,7 +1031,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
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: withhold the stream
|
||||
# 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
|
||||
|
||||
|
|
|
|||
|
|
@ -699,7 +699,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
``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 fails closed instead of releasing raw output.
|
||||
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
|
||||
|
|
@ -788,7 +789,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
# ------------------------------------------------------------------ #
|
||||
# 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 buffer truncated here still fails closed on text. #
|
||||
# delivered, so a truncated buffer still reports text undeliverable. #
|
||||
# ------------------------------------------------------------------ #
|
||||
if final_chunk.get("type") == "response.output_item.done":
|
||||
model_response_stream: Final = (
|
||||
|
|
@ -813,7 +814,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
# 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, and a #
|
||||
# rewrite a caller expects delivered fails closed instead. #
|
||||
# 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:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ 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 TYPE_CHECKING, Any, Final, Literal
|
||||
|
|
@ -77,7 +78,7 @@ class _StreamRewriteObserver(CustomGuardrail):
|
|||
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."""
|
||||
other translation are discarded by the executor, which releases the original chunks."""
|
||||
|
||||
def __init__(self, inner: CustomGuardrail) -> None:
|
||||
super().__init__(guardrail_name=inner.guardrail_name)
|
||||
|
|
@ -133,6 +134,19 @@ def _prepare_hook_input(
|
|||
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."""
|
||||
|
||||
|
|
@ -263,29 +277,37 @@ class PipelineExecutor:
|
|||
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."""
|
||||
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
|
||||
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,
|
||||
)
|
||||
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)
|
||||
return
|
||||
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
|
||||
raise UndeliverableStreamRewrite(step.guardrail)
|
||||
_release_original_chunks(step.guardrail, streaming_chunks, originals)
|
||||
|
||||
@staticmethod
|
||||
async def _run_step(
|
||||
|
|
@ -386,8 +408,6 @@ class PipelineExecutor:
|
|||
) # mutable-ok: modified-data contract is a plain dict
|
||||
return ("pass", response if isinstance(response, dict) else None, None, None)
|
||||
|
||||
except UndeliverableStreamRewrite:
|
||||
raise
|
||||
except Exception as e:
|
||||
if CustomGuardrail._is_guardrail_intervention(e):
|
||||
error_msg: Final = _extract_error_message(e)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from email.mime.text import MIMEText
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm import _custom_logger_compatible_callbacks_literal
|
||||
from litellm.constants import (
|
||||
|
|
@ -155,7 +155,7 @@ from litellm.proxy.hooks.sensitive_data_routing import (
|
|||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
|
||||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
|
|
@ -518,108 +518,72 @@ def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
|
|||
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
|
||||
|
||||
|
||||
class _PipelineErrorBody(TypedDict):
|
||||
message: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
policies: ReadOnly[tuple[str, ...]]
|
||||
guardrails: NotRequired[ReadOnly[tuple[str, ...]]]
|
||||
|
||||
|
||||
class _PipelineErrorDetail(TypedDict):
|
||||
error: ReadOnly[_PipelineErrorBody]
|
||||
|
||||
|
||||
def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) -> HTTPException:
|
||||
detail: Final[_PipelineErrorDetail] = {
|
||||
"error": {
|
||||
"message": (
|
||||
f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail "
|
||||
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,),
|
||||
"guardrails": (guardrail_name,),
|
||||
}
|
||||
}
|
||||
return HTTPException(status_code=400, detail=detail)
|
||||
|
||||
|
||||
def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
"""
|
||||
Reject up front the requests whose post_call pipelines could never run.
|
||||
|
||||
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
|
||||
(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
|
||||
if not is_stream and not is_background:
|
||||
return
|
||||
post_call_pipelines: Final = tuple(
|
||||
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 is released ungoverned by it: %s",
|
||||
policy_name,
|
||||
", ".join(unsupported),
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
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 the stream is
|
||||
released the way it was before pipelines ran on streams at all, with a
|
||||
warning naming what went ungoverned.
|
||||
"""
|
||||
post_call_pipelines: Final = _post_call_pipelines(request_data)
|
||||
if not post_call_pipelines:
|
||||
return
|
||||
post_call_policies: Final = tuple(policy_name for policy_name, _pipeline in post_call_pipelines)
|
||||
if is_background:
|
||||
background_detail: Final[_PipelineErrorDetail] = {
|
||||
"error": {
|
||||
"message": (
|
||||
"Policies with post_call guardrail pipelines cannot govern background "
|
||||
f"responses: {', '.join(post_call_policies)}. Retry with background=false."
|
||||
),
|
||||
"type": "guardrail_pipeline_error",
|
||||
"policies": post_call_policies,
|
||||
}
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=background_detail)
|
||||
step_guardrails: Final = tuple(
|
||||
dict.fromkeys(step.guardrail for _policy_name, pipeline in post_call_pipelines for step in pipeline.steps)
|
||||
)
|
||||
unsupported_guardrails: Final = tuple(
|
||||
guardrail for guardrail in step_guardrails if not _pipeline_step_supports_unified_streaming(guardrail)
|
||||
)
|
||||
if unsupported_guardrails:
|
||||
unsupported_detail: Final[_PipelineErrorDetail] = {
|
||||
"error": {
|
||||
"message": (
|
||||
"Policies with post_call guardrail pipelines cannot govern streaming responses "
|
||||
"because these pipeline guardrails do not support the unified apply_guardrail "
|
||||
f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or drop "
|
||||
"them from the pipeline steps so guardrails.add scans them on streamed output."
|
||||
),
|
||||
"type": "guardrail_pipeline_error",
|
||||
"policies": post_call_policies,
|
||||
"guardrails": unsupported_guardrails,
|
||||
}
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=unsupported_detail)
|
||||
return ()
|
||||
route: Final = user_api_key_dict.request_route
|
||||
if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None:
|
||||
return
|
||||
route_detail: Final[_PipelineErrorDetail] = {
|
||||
"error": {
|
||||
"message": (
|
||||
"Policies with post_call guardrail pipelines cannot govern streaming responses on "
|
||||
f"route {route} because it has no endpoint guardrail translation to scan the stream "
|
||||
f"through: {', '.join(post_call_policies)}. Retry with stream=false."
|
||||
),
|
||||
"type": "guardrail_pipeline_error",
|
||||
"policies": post_call_policies,
|
||||
}
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=route_detail)
|
||||
if route and resolve_endpoint_translation(user_api_key_dict, None) is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
|
||||
"(no endpoint guardrail translation); the stream is released ungoverned by them: %s",
|
||||
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)
|
||||
)
|
||||
|
||||
|
||||
def _prompt_block_text(block: object) -> str:
|
||||
|
|
@ -1990,7 +1954,7 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
try:
|
||||
_raise_for_streaming_post_call_pipelines(data, user_api_key_dict)
|
||||
_warn_background_skips_post_call_pipelines(data)
|
||||
|
||||
# Execute guardrail pipelines before the normal callback loop
|
||||
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
|
||||
|
|
@ -3371,11 +3335,7 @@ class ProxyLogging:
|
|||
1. /chat/completions
|
||||
"""
|
||||
caps: Final = ProxyLogging._callback_capabilities()
|
||||
post_call_pipelines: Final = tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in _policy_pipelines(request_data)
|
||||
if pipeline.mode == "post_call"
|
||||
)
|
||||
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
|
||||
|
|
@ -3486,9 +3446,11 @@ class ProxyLogging:
|
|||
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.
|
||||
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:
|
||||
|
|
@ -3498,39 +3460,27 @@ class ProxyLogging:
|
|||
|
||||
resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
|
||||
if resolved is None:
|
||||
policy_names: Final = tuple(policy_name for policy_name, _pipeline in pipelines)
|
||||
raise ProxyException(
|
||||
message=(
|
||||
"Policy pipelines could not govern this streaming response shape; "
|
||||
f"the response was withheld: {', '.join(policy_names)}."
|
||||
),
|
||||
type="guardrail_pipeline_error",
|
||||
param=None,
|
||||
code=500,
|
||||
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:
|
||||
try:
|
||||
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,
|
||||
)
|
||||
except UndeliverableStreamRewrite as rewrite:
|
||||
async for error_chunk in unified_guardrail.emit_streaming_http_error(
|
||||
_undeliverable_stream_rewrite_error(policy_name, rewrite.guardrail_name),
|
||||
call_type,
|
||||
buffered,
|
||||
request_data,
|
||||
):
|
||||
yield error_chunk
|
||||
return
|
||||
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
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Tests for the pipeline executor.
|
|||
Uses mock guardrails to validate pipeline execution without external services.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -941,7 +942,55 @@ class _TextTranslation:
|
|||
return responses_so_far
|
||||
|
||||
|
||||
async def _run_streaming_step(returned_texts, translation):
|
||||
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",
|
||||
|
|
@ -949,31 +998,41 @@ async def _run_streaming_step(returned_texts, translation):
|
|||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="p",
|
||||
streaming_chunks=[object()],
|
||||
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_rewrite_escapes_execute_steps_regardless_of_step_actions(monkeypatch):
|
||||
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 pytest.raises(UndeliverableStreamRewrite) as info:
|
||||
await _run_streaming_step(["hello [MASKED]"], translation)
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_streaming_step(translation, chunks)
|
||||
|
||||
assert info.value.guardrail_name == "masker"
|
||||
_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):
|
||||
async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch, caplog):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))])
|
||||
|
||||
result = await _run_streaming_step(("hello world",), _TextTranslation())
|
||||
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):
|
||||
|
|
@ -989,10 +1048,64 @@ class _InPlaceMutatingGuardrail(CustomGuardrail):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_in_place_rewrite_still_withholds_stream(monkeypatch):
|
||||
async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(monkeypatch, caplog):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_InPlaceMutatingGuardrail()])
|
||||
chunks = [_chunk()]
|
||||
|
||||
with pytest.raises(UndeliverableStreamRewrite) as info:
|
||||
await _run_streaming_step(["hello [MASKED]"], _TextTranslation())
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_streaming_step(_TextTranslation(), chunks)
|
||||
|
||||
assert info.value.guardrail_name == "masker"
|
||||
_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()]
|
||||
|
||||
|
||||
@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()]
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, List
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -24,9 +25,9 @@ from litellm.integrations.custom_guardrail import (
|
|||
ModifyResponseException,
|
||||
)
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import 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.utils import ProxyLogging, _streamable_post_call_pipelines
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
|
||||
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import (
|
||||
|
|
@ -1384,76 +1385,87 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once(
|
|||
assert seen["count"] == 1
|
||||
|
||||
|
||||
def _warnings(caplog: pytest.LogCaptureFixture) -> List[str]:
|
||||
return [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_verbatim(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = _stream_chunks()
|
||||
delivered: List[Any] = []
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
guardrails_only=True,
|
||||
)
|
||||
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,
|
||||
):
|
||||
delivered.append(item)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["policies"] == ("response-governance",)
|
||||
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
|
||||
assert "stream=false" in info.value.detail["error"]["message"]
|
||||
assert out is not None
|
||||
assert out.get("stream") is True
|
||||
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
|
||||
assert len(delivered) == 2
|
||||
assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
data = _post_call_pipeline_data(background=True)
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="aresponses",
|
||||
guardrails_only=True,
|
||||
)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["policies"] == ("response-governance",)
|
||||
assert "background=false" in info.value.detail["error"]["message"]
|
||||
assert out is not None
|
||||
assert out.get("background") is True
|
||||
assert any("response-governance" in message and "background" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(make_user_api_key_auth):
|
||||
post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
|
||||
auth = make_user_api_key_auth(request_route="/custom/stream")
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
data = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"background": True,
|
||||
"metadata": {
|
||||
"_guardrail_pipelines": [("request-governance", pre_call)],
|
||||
"_pipeline_managed_guardrails": {"gr-post"},
|
||||
},
|
||||
}
|
||||
|
||||
assert (
|
||||
_raise_for_streaming_post_call_pipelines(
|
||||
{"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="aresponses",
|
||||
guardrails_only=True,
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
_raise_for_streaming_post_call_pipelines(
|
||||
{"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
_raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
_raise_for_streaming_post_call_pipelines(
|
||||
{"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert _raise_for_streaming_post_call_pipelines({"stream": True}, auth) is None
|
||||
assert _raise_for_streaming_post_call_pipelines({"background": True}, auth) is None
|
||||
|
||||
assert out is not None
|
||||
assert not any("background" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -1485,6 +1497,56 @@ async def _async_chunk_iter(chunks: List[Any]):
|
|||
yield chunk
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported(
|
||||
make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
class NativeOnlyGuardrail(CustomGuardrail):
|
||||
pass
|
||||
|
||||
supported = _unified_stream_guardrail({})
|
||||
native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call)
|
||||
monkeypatch.setattr(litellm, "callbacks", [supported, native_only])
|
||||
governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
ungoverned = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")],
|
||||
)
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")])
|
||||
data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions"))
|
||||
|
||||
assert streamable == (("governed", governed),)
|
||||
assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog))
|
||||
assert not any("'governed'" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_is_empty_on_route_without_translation(
|
||||
make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail({})])
|
||||
governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
data = {"metadata": {"_guardrail_pipelines": [("governed", governed)]}}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/custom/stream"))
|
||||
|
||||
assert streamable == ()
|
||||
assert any("/custom/stream" in message and "governed" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_is_empty_without_post_call_pipelines(make_user_api_key_auth, caplog):
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
|
||||
auth = make_user_api_key_auth(request_route="/custom/stream")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
assert _streamable_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth) == ()
|
||||
assert _streamable_post_call_pipelines({"stream": True}, auth) == ()
|
||||
|
||||
assert _warnings(caplog) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"])
|
||||
async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified(
|
||||
|
|
@ -1507,21 +1569,25 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("native_lifecycle", [False, True])
|
||||
async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_unified_support(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle
|
||||
async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
if native_lifecycle:
|
||||
|
||||
class NativeOnlyGuardrail(CustomGuardrail):
|
||||
use_native_lifecycle_hooks = True
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
return inputs
|
||||
|
||||
else:
|
||||
|
||||
class NativeOnlyGuardrail(CustomGuardrail):
|
||||
pass
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
|
|
@ -1529,18 +1595,29 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni
|
|||
[NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
|
||||
)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = _stream_chunks()
|
||||
delivered: List[Any] = []
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
guardrails_only=True,
|
||||
)
|
||||
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,
|
||||
):
|
||||
delivered.append(item)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
|
||||
assert "apply_guardrail" in info.value.detail["error"]["message"]
|
||||
assert out is not None
|
||||
assert out.get("stream") is True
|
||||
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
|
||||
assert len(delivered) == 2
|
||||
assert seen.get("count") is None
|
||||
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1614,25 +1691,34 @@ async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
async def test_streaming_iterator_hook_releases_stream_when_route_has_no_guardrail_translation(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = _stream_chunks()
|
||||
delivered: List[Any] = []
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.pre_call_hook(
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
guardrails_only=True,
|
||||
)
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"),
|
||||
response=_async_chunk_iter(chunks),
|
||||
request_data=data,
|
||||
):
|
||||
delivered.append(item)
|
||||
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["policies"] == ("response-governance",)
|
||||
assert "/custom/stream" in info.value.detail["error"]["message"]
|
||||
assert out is not None
|
||||
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
|
||||
assert len(delivered) == 2
|
||||
assert seen.get("count") is None
|
||||
assert any("/custom/stream" in message and "response-governance" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1714,8 +1800,8 @@ 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")])
|
||||
async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewrite(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error
|
||||
async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog
|
||||
):
|
||||
transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731
|
||||
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
|
||||
|
|
@ -1724,7 +1810,7 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewr
|
|||
data = _post_call_pipeline_data(step=step, stream=True)
|
||||
delivered: List[Any] = []
|
||||
|
||||
async def _drain() -> None:
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
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(_tool_call_stream_chunks()),
|
||||
|
|
@ -1732,16 +1818,10 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewr
|
|||
):
|
||||
delivered.append(item)
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await _drain()
|
||||
|
||||
error = info.value.detail["error"]
|
||||
assert delivered == []
|
||||
assert info.value.status_code == 400
|
||||
assert error["type"] == "guardrail_pipeline_error"
|
||||
assert error["policies"] == ("response-governance",)
|
||||
assert error["guardrails"] == ("gr-post",)
|
||||
assert "stream=false" in error["message"]
|
||||
assert len(delivered) == 2
|
||||
assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}'
|
||||
assert delivered[1].choices[0].finish_reason == "tool_calls"
|
||||
assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1849,30 +1929,28 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = [object(), object()]
|
||||
delivered: List[Any] = []
|
||||
|
||||
async def _drain() -> None:
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
response=_async_chunk_iter([object(), object()]),
|
||||
response=_async_chunk_iter(chunks),
|
||||
request_data=data,
|
||||
):
|
||||
delivered.append(item)
|
||||
|
||||
with pytest.raises(ProxyException) as info:
|
||||
await _drain()
|
||||
|
||||
assert delivered == []
|
||||
assert info.value.code == "500"
|
||||
assert "withheld" in info.value.message
|
||||
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
|
||||
assert len(delivered) == 2
|
||||
assert seen.get("count") is None
|
||||
assert any("response-governance" in message and "shape" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def _anthropic_sse_chunks() -> List[bytes]:
|
||||
|
|
@ -1953,9 +2031,9 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_write_back(monkeypatch):
|
||||
async def test_pipeline_executor_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog):
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
|
||||
|
||||
class NoWriteBackTranslation(BaseTranslation):
|
||||
async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj):
|
||||
|
|
@ -1984,45 +2062,23 @@ async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_w
|
|||
|
||||
transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731
|
||||
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)])
|
||||
chunks = _stream_chunks()
|
||||
|
||||
with pytest.raises(UndeliverableStreamRewrite):
|
||||
await PipelineExecutor.execute_steps(
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = 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(),
|
||||
streaming_chunks=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
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
delivered: List[Any] = []
|
||||
|
||||
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(_stream_chunks()),
|
||||
request_data=data,
|
||||
):
|
||||
delivered.append(item)
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await _drain()
|
||||
|
||||
assert delivered == []
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["pipeline_context"]["step_results"] == [
|
||||
{"guardrail": "gr-post", "outcome": "error", "action": "block"}
|
||||
]
|
||||
assert result.terminal_action == "allow"
|
||||
assert [chunk.choices[0].delta.content for chunk in chunks] == ["hello ", "world"]
|
||||
assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue