diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c422a7c0964..acd2c2c973a 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,8 +6,11 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal +from pydantic import BaseModel + import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -24,8 +27,10 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStep, PipelineStepResult, ) +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, ) @@ -36,6 +41,90 @@ except ImportError: HTTPException = None +class UndeliverableStreamRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the streamed response, which streaming pipelines cannot deliver" + ) + self.guardrail_name: Final = guardrail_name + + +def _tool_call_shape(tool_call: object) -> tuple[object, object]: + plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + function: Final = plain.get("function") if isinstance(plain, Mapping) else None + if not isinstance(function, Mapping): + return (None, None) + return (function.get("name"), function.get("arguments")) + + +def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool: + return sent is not None and returned is not None and list(returned) != list(sent) + + +def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool: + if sent is None or returned is None: + return False + return [_tool_call_shape(tool_call) for tool_call in returned] != [ + _tool_call_shape(tool_call) for tool_call in sent + ] + + +class _StreamRewriteObserver(CustomGuardrail): + """Stand-in handed to the endpoint translation in place of a streaming pipeline step's + guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate + withholds the stream whenever the guardrail returned different output than it was given, + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime.""" + + def __init__(self, inner: CustomGuardrail) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.rewrote = False + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + outputs: Final = await self.inner.apply_guardrail( + inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj + ) + self.rewrote = ( + self.rewrote + or _rewrote_texts(inputs.get("texts"), outputs.get("texts")) + or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls")) + ) + return outputs + + +def _prepare_hook_input( + step: PipelineStep, + callback: CustomLogger, + data: dict, # mutable-ok: same request-payload shape the hooks mutate + raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data +) -> tuple[dict, bool]: # mutable-ok: returns that same request-payload dict + """Inject the step's guardrail name into metadata so should_run_guardrail() allows it, + and pick the payload the step scans: a scan_raw_request step evaluates the pristine + pre-pipeline snapshot instead of `data` (which earlier pass_data steps in this same + pipeline may have already rewritten), same reason the normal sequential/parallel + guardrail loops do this.""" + if "metadata" not in data: + data["metadata"] = {} + data["metadata"]["guardrails"] = [step.guardrail] + + scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data + independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None else data + ) + if hook_input is not data: + hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] + return hook_input, scans_raw_request + + class PipelineExecutor: """Executes guardrail pipelines with ordered, conditional step logic.""" @@ -195,23 +284,7 @@ class PipelineExecutor: return ("error", None, f"Guardrail '{step.guardrail}' not found", None) try: - # Inject guardrail name into metadata so should_run_guardrail() allows it - if "metadata" not in data: - data["metadata"] = {} - data["metadata"]["guardrails"] = [step.guardrail] - - # A scan_raw_request step evaluates the pristine pre-pipeline - # snapshot instead of `data` (which earlier pass_data steps in - # this same pipeline may have already rewritten), same reason - # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) - hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data - independent_snapshot(raw_request_snapshot) - if scans_raw_request and raw_request_snapshot is not None - else data - ) - if hook_input is not data: - hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] + hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) # Use unified_guardrail path if callback implements apply_guardrail target: CustomLogger = callback @@ -239,13 +312,16 @@ class PipelineExecutor: f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", None, ) + observer: Final = _StreamRewriteObserver(callback) await endpoint_translation.process_output_streaming_response( responses_so_far=streaming_chunks, - guardrail_to_apply=callback, + guardrail_to_apply=observer, litellm_logging_obj=data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, request_data=hook_input, ) + if observer.rewrote: + raise UndeliverableStreamRewrite(step.guardrail) response = None elif mode == "post_call": response = await target.async_post_call_success_hook( @@ -269,6 +345,8 @@ class PipelineExecutor: return ("pass", {"response": response}, None, None) 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 138f272af42..6c990222e51 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -155,7 +155,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at -from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -511,6 +511,23 @@ 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, and streaming pipelines cannot deliver " + "rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add " + "applies it to streamed output." + ), + "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. @@ -3468,11 +3485,12 @@ class ProxyLogging: pipeline allows it), then runs each pipeline's steps against the assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow - releases the buffered chunks as that machinery left them (the - Responses and A2A translations write guardrail output back into the - final chunk, exactly as they do for flat guardrails); a block or - modify_response terminates with the translation's block chunks or the - raised error. + releases the buffered chunks verbatim; a step whose guardrail rewrote + the output withholds the stream with a 400 instead, since no + translation rewrites every buffered chunk consistently and some + rewrites (Bedrock's ANONYMIZED action, for one) are only decided at + runtime; a block or modify_response terminates with the translation's + block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: @@ -3495,16 +3513,26 @@ class ProxyLogging: call_type, endpoint_translation = resolved for policy_name, pipeline in pipelines: - result: PipelineExecutionResult = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode="post_call", - data=request_data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - policy_name=policy_name, - streaming_chunks=buffered, - endpoint_translation=endpoint_translation, - ) + try: + 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 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 054a5af4148..52fd8777a19 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -13,7 +13,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( CustomCodeGuardrail, ) -from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -811,3 +811,64 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch): assert outcome == "pass" assert guardrail.native_pre_call_ran is True assert "guardrail_to_apply" not in data + + +class _TextReturningGuardrail(CustomGuardrail): + def __init__(self, returned_texts): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.returned_texts = returned_texts + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": self.returned_texts} + + +class _TextTranslation: + def __init__(self): + self.seen_guardrail_names = [] + + async def process_output_streaming_response( + self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + self.seen_guardrail_names.append(guardrail_to_apply.guardrail_name) + await guardrail_to_apply.apply_guardrail( + inputs={"texts": ["hello world"]}, + request_data=request_data or {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + +async def _run_streaming_step(returned_texts, translation): + return await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=[object()], + endpoint_translation=translation, + ) + + +@pytest.mark.asyncio +async def test_streaming_step_rewrite_escapes_execute_steps_regardless_of_step_actions(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + translation = _TextTranslation() + + with pytest.raises(UndeliverableStreamRewrite) as info: + await _run_streaming_step(["hello [MASKED]"], translation) + + assert info.value.guardrail_name == "masker" + assert translation.seen_guardrail_names == ["masker"] + + +@pytest.mark.asyncio +async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))]) + + result = await _run_streaming_step(("hello world",), _TextTranslation()) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] 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 895a986f348..74bd2e483cc 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,7 +11,7 @@ from __future__ import annotations import asyncio import json -from typing import Any, Dict, List +from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -878,10 +878,12 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p # --------------------------------------------------------------------------- -def _post_call_pipeline_data(guardrail: str = "gr-post", **extra: Any) -> Dict[str, Any]: +def _post_call_pipeline_data( + guardrail: str = "gr-post", step: PipelineStep | None = None, **extra: Any +) -> Dict[str, Any]: pipeline = GuardrailPipeline( mode="post_call", - steps=[PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], + steps=[step or PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], ) return { "model": "m", @@ -1587,6 +1589,101 @@ async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks( assert "output blocked" in str(info.value.detail) +def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, Any]]) -> CustomGuardrail: + class RewritingStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, **transform(inputs)} + + return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _tool_call_stream_chunks() -> List[Any]: + tool_call = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"ssn": "123"}'}, + } + return [ + litellm.ModelResponseStream( + choices=[{"index": 0, "delta": {"tool_calls": [tool_call]}, "finish_reason": None}] + ), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]), + ] + + +def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: + return [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) +@pytest.mark.parametrize( + "make_chunks, transform", + [ + (_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}), + (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}), + ], + ids=["texts", "tool_calls"], +) +async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error) + data = _post_call_pipeline_data(step=step, 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(make_chunks()), + request_data=data, + ): + 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"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_chunks, transform", + [ + (_stream_chunks, lambda inputs: {"texts": tuple(inputs["texts"])}), + (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "123"}')}), + ], + ids=["texts_as_tuple", "tool_calls_as_dicts"], +) +async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_another_shape( + proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = make_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape( proxy_logging, make_user_api_key_auth, monkeypatch