diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8dc6881d23e..a6e3d000120 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -762,6 +762,9 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + def rewrites_streamed_output(self) -> bool: + return self.mask_response_content + def _deployment_pre_call_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 722f96ef814..85eb50c78e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1947,6 +1947,15 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str=exception_str, ) + def rewrites_streamed_output(self) -> bool: + return ( + super().rewrites_streamed_output() + or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns) + or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values()) + or any(action == ContentFilterAction.MASK for _, _, action in self.category_keywords.values()) + or any(action == ContentFilterAction.MASK for _, _, action in self.always_block_category_keywords.values()) + ) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..89527c05eb6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -19,7 +19,7 @@ from litellm.cost_calculator import _infer_call_type from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import load_guardrail_translation_mappings +from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -62,6 +62,36 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran return translation +def resolve_endpoint_translation( + user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None +) -> "tuple[str, BaseTranslation] | None": + """ + Resolve the endpoint guardrail translation for a streamed response: the + request route wins, falling back to inferring the call type from the first + response chunk (the same resolution order the streaming iterator hook uses). + Returns None when the call type is unresolvable or has no translation. + """ + route_call_types: Final = ( + get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None + ) + call_type: Final = ( + route_call_types[0].value + if route_call_types + else ( + _infer_call_type(call_type=None, completion_response=first_response_item) + if first_response_item is not None + else None + ) + ) + if call_type is None: + return None + try: + handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type)) + except ValueError: + return None + return call_type, handler_cls() + + def _chunk_choices(item: object) -> Sequence[object]: choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] return choices @@ -346,7 +376,7 @@ class UnifiedLLMGuardrails(CustomLogger): return response - async def _handle_streaming_block( + async def handle_streaming_block( self, exc: "ModifyResponseException", endpoint_translation: _EndpointTranslation, @@ -402,7 +432,7 @@ class UnifiedLLMGuardrails(CustomLogger): return None return call_type - async def _emit_streaming_http_error( + async def emit_streaming_http_error( self, exc: HTTPException, call_type: str | None, @@ -577,7 +607,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -586,7 +616,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + async for error_item in self.emit_streaming_http_error(e, call_type, responses_so_far, request_data): yield error_item raise _StreamTerminated() @@ -758,7 +788,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -846,6 +876,14 @@ class UnifiedLLMGuardrails(CustomLogger): choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) + def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object: + """Streaming flag resolution order (later wins): default < guardrail + attribute < guardrail_config dict < this callback's optional_params.""" + attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default) + config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None) + config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value + return self.optional_params.get(name, config_value) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -876,17 +914,8 @@ class UnifiedLLMGuardrails(CustomLogger): if guardrail_to_apply is None: guardrail_to_apply = request_data.pop("guardrail_to_apply", None) - # Get streaming configuration. Resolution order (later wins): default - # < guardrail attribute < guardrail_config dict < this callback's - # optional_params. def _streaming_flag(name: str, default: object) -> Any: - value = default - if guardrail_to_apply is not None: - value = getattr(guardrail_to_apply, name, value) - config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) - if isinstance(config, dict): - value = config.get(name, value) - return self.optional_params.get(name, value) + return self.resolve_streaming_flag(guardrail_to_apply, name, default) sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). @@ -1060,7 +1089,7 @@ class UnifiedLLMGuardrails(CustomLogger): # The current chunk was appended to responses_so_far but not # yet yielded, so exclude it: the continuation must reflect # only what the client has actually received. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=chunks_yielded, @@ -1124,7 +1153,7 @@ class UnifiedLLMGuardrails(CustomLogger): # terminating SSE sequence with the block message rather than # propagating into a bare error blob that truncates the stream. # The withheld original chunks are never released. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 190784a2a60..acd2c2c973a 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,7 +6,10 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time -from typing import Any, Final, Literal +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 @@ -24,6 +27,13 @@ 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, + ) try: from fastapi.exceptions import HTTPException @@ -31,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.""" @@ -43,6 +137,8 @@ class PipelineExecutor: call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ Execute pipeline steps sequentially with conditional actions. @@ -59,6 +155,12 @@ class PipelineExecutor: step whose guardrail opted into ``scan_raw_request`` evaluates the original request instead of whatever an earlier ``pass_data`` step in this same pipeline already rewrote. + streaming_chunks: buffered chunks of a completed stream. When set + (with ``endpoint_translation``), post_call steps scan the + assembled streamed output through the endpoint translation + instead of calling ``async_post_call_success_hook``. + endpoint_translation: the guardrail translation for the streamed + endpoint, resolved by the caller. Returns: PipelineExecutionResult with terminal action and step results @@ -83,6 +185,8 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, call_type=call_type, raw_request_snapshot=raw_request_snapshot, + streaming_chunks=streaming_chunks, + endpoint_translation=endpoint_translation, ) duration = time.perf_counter() - start_time @@ -154,6 +258,8 @@ class PipelineExecutor: user_api_key_dict: Any, call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], dict | None, @@ -178,30 +284,12 @@ 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 - use_unified: Final = ( - "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks - ) - if use_unified: + use_unified: Final = PipelineExecutor.supports_unified_execution(callback) + if use_unified and streaming_chunks is None: hook_input["guardrail_to_apply"] = callback target = UnifiedLLMGuardrails() @@ -216,6 +304,25 @@ class PipelineExecutor: callback.mark_pre_call_hook_ran(data) if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) + elif mode == "post_call" and streaming_chunks is not None: + if not use_unified or endpoint_translation is None: + return ( + "error", + None, + f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", + None, + ) + observer: Final = _StreamRewriteObserver(callback) + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=data.get("litellm_logging_obj"), + 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( user_api_key_dict=user_api_key_dict, @@ -238,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) @@ -246,6 +355,12 @@ class PipelineExecutor: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + @staticmethod + def supports_unified_execution(callback: CustomGuardrail) -> bool: + """Whether this guardrail runs through the unified apply_guardrail path, + the interface streaming pipeline execution requires.""" + return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: """Look up an initialized guardrail callback by name from litellm.callbacks.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 80e991640a9..6c990222e51 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 ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -139,6 +139,7 @@ from litellm.proxy.db.token_auth import ( ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, + resolve_endpoint_translation, ) from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck @@ -154,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 ( @@ -486,29 +487,139 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("stream") is not True and data.get("background") is not True: +def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + return callback is not None and PipelineExecutor.supports_unified_execution(callback) + + +def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool: + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + if callback is None: + return False + transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only") + return callback.rewrites_streamed_output() or transform_mode == "incremental_diff" + + +class _PipelineErrorBody(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + 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, 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. + + Background responses skip the post_call hooks entirely, so a pipeline + governing one would silently never execute. Streaming responses execute + pipelines against the buffered stream through the endpoint guardrail + translation of the request route, releasing the buffered chunks on allow. + That needs every step's guardrail to support the unified apply_guardrail + interface and to only allow or block (a step that rewrites streamed + content, via mask_response_content, a MASK action, or + streaming_transform_mode=incremental_diff, would have its rewrite silently + dropped), and needs the route to have a translation at all; anything else + keeps the 400 rather than letting ungoverned output stream through. + """ + 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_policies: Final = tuple( - policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" + post_call_pipelines: Final = tuple( + (policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" ) - if not post_call_policies: + if not post_call_pipelines: return - raise HTTPException( - status_code=400, - detail={ + 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 streaming or background " - f"responses yet: {', '.join(post_call_policies)}. Retry with stream=false and " - "background=false, or move these policies' output guardrails from pipeline steps to " - "guardrails.add, which scans streamed output." + "Policies with post_call guardrail pipelines cannot govern background " + f"responses: {', '.join(post_call_policies)}. Retry with background=false." ), "type": "guardrail_pipeline_error", - "policies": list(post_call_policies), + "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) + rewriting_guardrails: Final = tuple( + guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail) + ) + if rewriting_guardrails: + rewriting_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses " + "because these pipeline guardrails rewrite streamed content (mask_response_content, " + "a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release " + f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop " + "them from the pipeline steps so guardrails.add applies them to streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": post_call_policies, + "guardrails": rewriting_guardrails, + } + } + raise HTTPException(status_code=400, detail=rewriting_detail) + route: Final = user_api_key_dict.request_route + if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None: + return + 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) def _prompt_block_text(block: object) -> str: @@ -1689,7 +1800,7 @@ class ProxyLogging: result: PipelineExecutionResult, data: dict, policy_name: str, - original_response: LLMResponseTypes | None = None, + original_response: "LLMResponseTypes | Sequence[object] | None" = None, ) -> dict: """ Handle a PipelineExecutionResult — allow, block, or modify_response. @@ -1699,7 +1810,9 @@ class ProxyLogging: payload (already sent upstream) must stay untouched; a replacement response carried in ``modified_data`` is adopted by the caller, and metadata-bucket writes (applied guardrails, guardrail logging info) - are merged back so headers and spend logs still see them. + are merged back so headers and spend logs still see them. On the + streaming path it is the buffered chunk list, carried into + ``ModifyResponseException.original_response`` for usage reporting. """ if result.terminal_action == "allow": if result.modified_data is not None: @@ -1862,7 +1975,7 @@ class ProxyLogging: ) try: - _raise_for_streaming_post_call_pipelines(data) + _raise_for_streaming_post_call_pipelines(data, user_api_key_dict) # Execute guardrail pipelines before the normal callback loop data, _ = await self._maybe_execute_pipelines( @@ -3195,11 +3308,16 @@ class ProxyLogging: # dict lookups + llm_router.get_deployment() per callback per chunk. _cached_guardrail_data: dict | None = None _guardrail_data_computed = False + pipeline_managed: Final = ( + _pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() + ) for callback in litellm.callbacks: try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): + if callback.guardrail_name in pipeline_managed: + continue # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks @@ -3256,12 +3374,17 @@ 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" + ) # Fast path: no real overrides. Internal proxy CustomLogger callbacks # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default # ``async for chunk: yield chunk`` body, so wrapping the iterator # through each of them adds N pass-through trampolines per chunk for # zero behavior change. Skip the chain entirely and stream through. - if not caps.iterator_overrides: + if not caps.iterator_overrides and not post_call_pipelines: try: async for chunk in response: yield chunk @@ -3281,8 +3404,11 @@ class ProxyLogging: current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) + pipeline_managed_names: Final = _pipeline_managed_guardrail_names(request_data, "post_call") for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): + if resolved_callback.guardrail_name in pipeline_managed_names: + continue if ( resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True @@ -3322,6 +3448,14 @@ class ProxyLogging: ), ) + if post_call_pipelines: + current_response = self._pipeline_gated_stream( + response=current_response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + pipelines=post_call_pipelines, + ) + try: async for chunk in current_response: yield chunk @@ -3337,6 +3471,90 @@ class ProxyLogging: # we reach this point the metadata is fully populated. ProxyLogging._fire_deferred_stream_logging(request_data) + async def _pipeline_gated_stream( + self, + response: "AsyncGenerator[object, None]", + user_api_key_dict: UserAPIKeyAuth, + request_data: dict, + pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + ) -> "AsyncGenerator[Any, None]": + """ + Execute post_call policy pipelines against a streamed response. + + Buffers the whole stream (nothing reaches the client until every + pipeline allows it), then runs each pipeline's steps against the + assembled output through the endpoint guardrail translation, the same + machinery flat post_call guardrails use at end of stream. An allow + releases the buffered chunks verbatim; 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: + buffered.append(item) + if not buffered: + return + + 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, + ) + 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 + try: + ProxyLogging._handle_pipeline_result( + result, data=request_data, policy_name=policy_name, original_response=buffered + ) + except ModifyResponseException as e: + if e.original_response is None: + e.original_response = buffered + async for block_chunk in unified_guardrail.handle_streaming_block( + e, endpoint_translation, stream_started=False, responses_so_far=() + ): + yield block_chunk + return + except HTTPException as e: + async for error_chunk in unified_guardrail.emit_streaming_http_error( + e, call_type, buffered, request_data + ): + yield error_chunk + return + + for buffered_item in buffered: + yield buffered_item + @staticmethod def _fire_deferred_stream_logging(request_data: dict) -> None: """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index be55ac47bde..73020fe3e6f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -3068,3 +3068,59 @@ class TestContentFilterToolCallArguments: request_data={}, input_type="response", ) + + +class TestRewritesStreamedOutput: + def test_block_only_rules_do_not_rewrite(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.BLOCK)], + blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], + ) + + assert guardrail.rewrites_streamed_output() is False + + def test_mask_blocked_word_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + blocked_words=[BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK)], + ) + + assert guardrail.rewrites_streamed_output() is True + + def test_mask_pattern_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)], + ) + + assert guardrail.rewrites_streamed_output() is True + + def test_mask_response_content_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], + mask_response_content=True, + ) + + assert guardrail.rewrites_streamed_output() is True + + @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) + def test_category_keywords_follow_the_category_action(self, action, expected): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + categories=[{"category": "bias_gender", "enabled": True, "action": action}], + ) + + assert guardrail.category_keywords and not guardrail.always_block_category_keywords + assert guardrail.rewrites_streamed_output() is expected + + @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) + def test_always_block_category_keywords_follow_the_category_action(self, action, expected): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + categories=[{"category": "age_discrimination", "enabled": True, "action": action}], + ) + + assert guardrail.always_block_category_keywords and not guardrail.category_keywords + assert guardrail.rewrites_streamed_output() is expected diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..fe4acbf8277 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1026,7 +1026,7 @@ class TestStreamingTransform: ) emitted = [] - async for item in handler._emit_streaming_http_error( + async for item in handler.emit_streaming_http_error( exc, call_type=CallTypes.asend_message.value, responses_so_far=[{"id": "req-1"}], 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 34a25d75bc0..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 @@ -10,7 +10,8 @@ Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, from __future__ import annotations import asyncio -from typing import Any, Dict, List +import json +from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -23,9 +24,11 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import ProxyException 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.types.guardrails import GuardrailEventHooks +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 ( GuardrailPipeline, PipelineStep, @@ -875,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", @@ -1284,7 +1289,8 @@ async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( ) assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ["response-governance"] + 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"] @@ -1304,32 +1310,511 @@ async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( ) assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ["response-governance"] + assert info.value.detail["error"]["policies"] == ("response-governance",) assert "background=false" in info.value.detail["error"]["message"] -def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(): +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") assert ( _raise_for_streaming_post_call_pipelines( - {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth ) is None ) assert ( _raise_for_streaming_post_call_pipelines( - {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth ) is None ) - assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) 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)]}} + {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth ) is None ) - assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None - assert _raise_for_streaming_post_call_pipelines({"background": True}) 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 + + +# --------------------------------------------------------------------------- +# post_call pipelines on streaming responses +# --------------------------------------------------------------------------- + + +def _unified_stream_guardrail(seen: Dict[str, Any], block: bool = False) -> CustomGuardrail: + class UnifiedStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + seen["count"] = seen.get("count", 0) + 1 + seen["input_type"] = input_type + if block: + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + return inputs + + return UnifiedStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _stream_chunks() -> List[Any]: + return [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "hello "}, "finish_reason": None}]), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}, "finish_reason": "stop"}]), + ] + + +async def _async_chunk_iter(chunks: List[Any]): + for chunk in chunks: + yield chunk + + +@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( + proxy_logging, make_user_api_key_auth, monkeypatch, request_route +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route=request_route), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert out is not None + assert out.get("stream") is True + + +@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 +): + if native_lifecycle: + + class NativeOnlyGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + else: + + class NativeOnlyGuardrail(CustomGuardrail): + pass + + monkeypatch.setattr( + litellm, + "callbacks", + [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "apply_guardrail" in info.value.detail["error"]["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rewrite_attribute, value", + [ + ("mask_response_content", True), + ("streaming_transform_mode", "incremental_diff"), + ("guardrail_config", {"streaming_transform_mode": "incremental_diff"}), + ], +) +async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content( + proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value +): + seen: Dict[str, Any] = {} + guardrail = _unified_stream_guardrail(seen) + setattr(guardrail, rewrite_attribute, value) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "rewrite streamed content" in info.value.detail["error"]["message"] + assert seen.get("count") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)]) +async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks( + proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="persimmon", action=action)], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + if not rejected: + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + assert out is not None and out.get("stream") is True + return + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "a MASK action" in info.value.detail["error"]["message"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + categories=[{"category": "bias_gender", "enabled": True, "action": "MASK"}], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + + +@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 +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["policies"] == ("response-governance",) + assert "/custom/stream" in info.value.detail["error"]["message"] + assert seen.get("count") is None + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_allow_releases_buffered_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + 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 = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert seen["count"] == 1 + assert seen["input_type"] == "response" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + 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 "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 +): + 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) + 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(), + response=_async_chunk_iter([object(), object()]), + 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 seen.get("count") is None + + +def _anthropic_sse_chunks() -> List[bytes]: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated_block( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep( + guardrail="gr-post", + on_pass="allow", + on_fail="modify_response", + modify_response_message="content policy block", + ) + ], + ) + data = _post_call_pipeline_data(stream=True) + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks = _anthropic_sse_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert seen["count"] == 1 + assert "content policy block" in raw + assert "hello world" not in raw + assert not any(item is chunk for item in delivered for chunk in chunks) + + +@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"} + ] + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + managed = RecordingGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + free = RecordingGuardrail( + guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [managed, free]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + ) + + assert result is not None + assert seen.get("gr-post") is None + assert seen["gr-free"] == 1