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 60b4444e1d4..89527c05eb6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -876,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, @@ -906,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). diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5642eb5383e..8b2c38d457f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 @@ -486,11 +487,19 @@ 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 _pipeline_step_supports_streaming(guardrail_name: str) -> bool: +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.mask_response_content or transform_mode == "incremental_diff" + + class _PipelineErrorBody(TypedDict): message: ReadOnly[str] type: ReadOnly[str] @@ -502,17 +511,20 @@ class _PipelineErrorDetail(TypedDict): error: ReadOnly[_PipelineErrorBody] -def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: +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 - translations, which requires every step's guardrail to support the unified - apply_guardrail interface; steps that cannot (native-lifecycle guardrails, - or guardrails not registered at all) keep the 400 rather than letting - ungoverned output stream through. + translation of the request route, releasing the buffered chunks verbatim + 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 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 @@ -536,30 +548,61 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None } } raise HTTPException(status_code=400, detail=background_detail) - unsupported_guardrails: Final = tuple( - dict.fromkeys( - step.guardrail - for _policy_name, pipeline in post_call_pipelines - for step in pipeline.steps - if not _pipeline_step_supports_streaming(step.guardrail) - ) + step_guardrails: Final = tuple( + dict.fromkeys(step.guardrail for _policy_name, pipeline in post_call_pipelines for step in pipeline.steps) ) - if not unsupported_guardrails: + 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 " + "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 - unsupported_detail: Final[_PipelineErrorDetail] = { + route_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 move " - "them from pipeline steps to guardrails.add, which scans streamed output." + "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, - "guardrails": unsupported_guardrails, } } - raise HTTPException(status_code=400, detail=unsupported_detail) + raise HTTPException(status_code=400, detail=route_detail) def _prompt_block_text(block: object) -> str: @@ -1915,7 +1958,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( @@ -3431,10 +3474,6 @@ class ProxyLogging: releases the buffered chunks verbatim; a block or modify_response terminates with the translation's block chunks or the raised error. """ - from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( - resolve_endpoint_translation, - ) - buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: buffered.append(item) @@ -3444,17 +3483,15 @@ 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) - unresolvable_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policy pipelines could not govern this streaming response shape; " - f"the response was withheld: {', '.join(policy_names)}." - ), - "type": "guardrail_pipeline_error", - "policies": policy_names, - } - } - raise HTTPException(status_code=500, detail=unresolvable_detail) + 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: 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 2b36ae111dc..6c90637158e 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 @@ -23,6 +23,7 @@ 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 @@ -1309,31 +1310,35 @@ async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( 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 # --------------------------------------------------------------------------- @@ -1366,15 +1371,16 @@ async def _async_chunk_iter(chunks: List[Any]): @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 + 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(), + user_api_key_dict=make_user_api_key_auth(request_route=request_route), data=data, call_type="completion", guardrails_only=True, @@ -1422,6 +1428,60 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni 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 +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 @@ -1490,10 +1550,10 @@ async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_ ): delivered.append(item) - with pytest.raises(HTTPException) as info: + with pytest.raises(ProxyException) as info: await _drain() assert delivered == [] - assert info.value.status_code == 500 - assert "withheld" in info.value.detail["error"]["message"] + assert info.value.code == "500" + assert "withheld" in info.value.message assert seen.get("count") is None