feat(guardrails): expose streaming knobs on generic_guardrail_api

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Optional
 
 from litellm.types.guardrails import SupportedGuardrailIntegrations
 
@@ -8,9 +8,23 @@
     from litellm.types.guardrails import Guardrail, LitellmParams
 
 
+def _get_config_value(litellm_params: Any, optional_params: Any, attribute_name: str) -> Optional[Any]:
+    if optional_params is not None:
+        value = (
+            optional_params.get(attribute_name)
+            if isinstance(optional_params, dict)
+            else getattr(optional_params, attribute_name, None)
+        )
+        if value is not None:
+            return value
+    return getattr(litellm_params, attribute_name, None)
+
+
 def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
     import litellm
 
+    optional_params = getattr(litellm_params, "optional_params", None)
+
     _generic_guardrail_api_callback = GenericGuardrailAPI(
         api_base=litellm_params.api_base,
         api_key=litellm_params.api_key,
@@ -22,6 +36,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         guardrail_name=guardrail.get("guardrail_name", ""),
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
+        streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"),
+        streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"),
     )
 
     litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback)

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -33,6 +33,7 @@
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
 
 GUARDRAIL_NAME = "generic_guardrail_api"
 
@@ -178,6 +179,8 @@ def __init__(
         unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
         fail_on_error: Optional[bool] = True,
         extra_headers: Optional[list] = None,
+        streaming_end_of_stream_only: Optional[bool] = None,
+        streaming_sampling_rate: Optional[int] = None,
         **kwargs,
     ):
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
@@ -209,6 +212,15 @@ def __init__(
 
         self.fail_on_error: bool = True if fail_on_error is None else fail_on_error
 
+        # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook
+        # via getattr(guardrail_to_apply, "streaming_*", default).
+        self.streaming_end_of_stream_only: bool = (
+            False if streaming_end_of_stream_only is None else streaming_end_of_stream_only
+        )
+        if streaming_sampling_rate is not None and streaming_sampling_rate < 1:
+            raise ValueError(f"streaming_sampling_rate must be >= 1 (got {streaming_sampling_rate})")
+        self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate
+
         # Set supported event hooks
         if "supported_event_hooks" not in kwargs:
             kwargs["supported_event_hooks"] = [
@@ -470,3 +482,11 @@ async def apply_guardrail(
             return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj)
         except Exception as e:
             return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj, is_unreachable=False)
+
+    @staticmethod
+    def get_config_model() -> Optional[type["GuardrailConfigModel"]]:
+        from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
+            GenericGuardrailAPIConfigModel,
+        )
+
+        return GenericGuardrailAPIConfigModel

litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py

@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field
-from typing_extensions import TYPE_CHECKING, TypedDict
+from typing_extensions import TypedDict
 
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -60,6 +60,30 @@ class GenericGuardrailAPIOptionalParams(BaseModel):
         ),
     )
 
+    streaming_end_of_stream_only: Optional[bool] = Field(
+        default=None,
+        description=(
+            "If False (default when unset), the guardrail runs on sampled chunks during "
+            "the stream at the cadence set by streaming_sampling_rate, and an in-flight "
+            "BLOCKED stops further chunks from streaming. If True, the guardrail runs "
+            "once at end of stream over the assembled response; lower cost and latency, "
+            "but flagged content has already streamed to the client before the terminal "
+            "block. Defaults are applied in GenericGuardrailAPI.__init__ when None so "
+            "unset optional_params does not shadow top-level litellm_params."
+        ),
+    )
+
+    streaming_sampling_rate: Optional[int] = Field(
+        default=None,
+        ge=1,
+        description=(
+            "When streaming_end_of_stream_only is False, the guardrail runs every Nth "
+            "streamed chunk. Ignored when streaming_end_of_stream_only is True. "
+            "Must be >= 1 when set. Defaults to 5 in GenericGuardrailAPI.__init__ "
+            "when None so unset optional_params does not shadow top-level litellm_params."
+        ),
+    )
+
 
 class GenericGuardrailAPIConfigModel(
     GuardrailConfigModel[GenericGuardrailAPIOptionalParams],