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 d4d23cd2e37..049900dbc46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,7 +8,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Mapping, Optional, Union from fastapi import HTTPException @@ -804,6 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, + streaming_flag_defaults: Optional[Mapping[str, Any]] = None, ) -> AsyncGenerator[Any, None]: """ Passes the entire stream to the guardrail @@ -827,10 +828,10 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply: CustomGuardrail = 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. + # < caller-supplied default < guardrail attribute < guardrail_config + # dict < this callback's optional_params. def _streaming_flag(name: str, default: Any) -> Any: - value = default + value = default if streaming_flag_defaults is None else streaming_flag_defaults.get(name, default) if guardrail_to_apply is not None: value = getattr(guardrail_to_apply, name, value) config = getattr(guardrail_to_apply, "guardrail_config", {}) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e6e864a8e13..eb1570a3332 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2705,6 +2705,10 @@ class ProxyLogging: user_api_key_dict=user_api_key_dict, request_data=request_data, response=current_response, + # A guardrail that ships its own iterator hook scans + # every chunk there, so keep that cadence when we + # reroute it: sampling would forward unscanned content. + streaming_flag_defaults=({"streaming_sampling_rate": 1} if effective_kind != kind else None), ), ) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 82b892f1923..2c4ea85d818 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -485,9 +485,6 @@ def _content_filter_guardrail(): blocked_words=[{"keyword": "zebra", "action": "BLOCK"}], default_on=True, ) - # scan every chunk so the block decision lands before the offending delta - # is forwarded, matching the /chat/completions withholding behavior - guardrail.streaming_sampling_rate = 1 return guardrail