diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index da30d85e26b..e486be12fe2 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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 diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 770fac6e443..afd8e0f67f7 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -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( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 4b7b1cc2700..80292aef2cf 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -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 diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b543c7f1e17..b0f79552bc5 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -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: diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 160938b723c..970ac20487c 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -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) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ca8e48b5dfa..d7bf832bca4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index d399b4f01cb..0685bc6aa1e 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -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()] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 2e73bebb07c..ad2b4a0efbb 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -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