feat(guardrails): per-mode stream_scope with bedrock stream and pass-through fixes

litellm/constants.py

@@ -1,10 +1,20 @@
 import os
 import sys
+from enum import Enum
 from types import MappingProxyType
 from typing import Final, Literal
 
 from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none
 
+SERVER_STREAMING_CLASSIFICATION_KEY: Final = "litellm_server_streaming_classification"
+
+
+class ServerStreamingClassification(str, Enum):
+    MARKER = "litellm-server-streaming"
+
+
+SERVER_STREAMING_CLASSIFICATION_MARKER: Final = ServerStreamingClassification.MARKER
+
 DEFER_PYDANTIC_BUILD: Final = os.getenv("DEFER_PYDANTIC_BUILD", "true") in ("true", "1", "on")
 DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
 AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))

litellm/integrations/custom_guardrail.py

@@ -6,7 +6,7 @@
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
 
 import httpx
 
@@ -21,11 +21,14 @@
 )
 from litellm.secret_managers.main import str_to_bool
 from litellm.types.guardrails import (
+    DEFAULT_GUARDRAIL_STREAM_SCOPE,
     DynamicGuardrailParams,
     GuardrailEventHooks,
+    GuardrailStreamScope,
     LitellmParams,
     LoggingOnlyScope,
     Mode,
+    runtime_stream_scope,
 )
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
@@ -49,6 +52,8 @@
     GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS,
     LOGS_GUARDRAIL_INFORMATION_MARKER,
     PRE_CALL_EXECUTED_GUARDRAILS_KEY,
+    SERVER_STREAMING_CLASSIFICATION_KEY,
+    SERVER_STREAMING_CLASSIFICATION_MARKER,
 )
 from litellm.exceptions import (
     BlockedPiiEntityError,
@@ -173,6 +178,42 @@ def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None
     return None
 
 
+_REALTIME_STREAMING_HOOKS: Final = frozenset({GuardrailEventHooks.realtime_input_transcription})
+
+
+def without_server_streaming_classification(data: Mapping[str, object]) -> dict[str, object]:
+    return {
+        key: value
+        for key, value in data.items()
+        if key != SERVER_STREAMING_CLASSIFICATION_KEY or value != SERVER_STREAMING_CLASSIFICATION_MARKER
+    }
+
+
+def guardrail_request_data_with_streaming(
+    data: Mapping[str, object],
+    *,
+    is_streaming: bool,
+) -> dict[str, object]:
+    data_without_server_classification: Final = without_server_streaming_classification(data)
+    if not is_streaming:
+        return data_without_server_classification
+    return {
+        **data_without_server_classification,
+        SERVER_STREAMING_CLASSIFICATION_KEY: SERVER_STREAMING_CLASSIFICATION_MARKER,
+    }
+
+
+def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool:
+    if event_type in _REALTIME_STREAMING_HOOKS:
+        return True
+    if not isinstance(data, Mapping):
+        return False
+    return (
+        data.get("stream") is True
+        or data.get(SERVER_STREAMING_CLASSIFICATION_KEY) is SERVER_STREAMING_CLASSIFICATION_MARKER
+    )
+
+
 class CustomGuardrail(CustomLogger):
     # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
     use_native_during_call_hook: ClassVar[bool] = False
@@ -183,6 +224,9 @@ class CustomGuardrail(CustomLogger):
     records_own_guardrail_information: ClassVar[bool] = False
     logging_only_scope: LoggingOnlyScope | None
 
+    stream_scope_default: GuardrailStreamScope = DEFAULT_GUARDRAIL_STREAM_SCOPE
+    stream_scope_by_hook: tuple[tuple[str, GuardrailStreamScope], ...] = ()
+
     timeout: float | httpx.Timeout | None = None
 
     def __init_subclass__(cls, **kwargs: object) -> None:  # kwargs-ok: forwarded to cooperative __init_subclass__ hooks
@@ -258,6 +302,8 @@ def __init__(
         self.run_in_parallel: bool = run_in_parallel
         self.scan_raw_request: bool = scan_raw_request
         self.only_scan_new_messages: bool = only_scan_new_messages
+        stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None))  # cast-ok: config
+        self.apply_stream_scope(stream_scope_arg)
         self.logging_only_scope = None
         if timeout is not None:
             self.timeout = timeout
@@ -1099,6 +1145,23 @@ def _suppressed_by_auto_router_compression(self) -> bool:
 
         return name in suppressed_compression_guardrails()
 
+    def apply_stream_scope(self, stream_scope: object) -> None:
+        default, by_hook = runtime_stream_scope(stream_scope)
+        self.stream_scope_default = default
+        self.stream_scope_by_hook = tuple(by_hook.items())
+
+    def stream_scope_allows(self, data: object, event_type: GuardrailEventHooks) -> bool:
+        scope: Final = next(
+            (scope for hook, scope in self.stream_scope_by_hook if hook == event_type.value),
+            self.stream_scope_default,
+        )
+        if scope == "both":
+            return True
+        is_streaming: Final = _request_is_streaming(data, event_type)
+        if scope == "streaming":
+            return is_streaming
+        return not is_streaming
+
     def should_run_guardrail(
         self,
         data,
@@ -1142,8 +1205,10 @@ def should_run_guardrail(
                         data, self.event_hook, event_type
                     )
                     if result is not None:
-                        return result
-                return True
+                        tagged_result: Final[bool] = bool(cast(object, result))  # cast-ok: helper return
+                        data_obj: Final[object] = cast(object, data)  # cast-ok: data param
+                        return tagged_result and self.stream_scope_allows(data_obj, event_type)
+                return self.stream_scope_allows(cast(object, data), event_type)  # cast-ok: data param
             return False
 
         if (
@@ -1167,8 +1232,9 @@ def should_run_guardrail(
                 )
             result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag(data, self.event_hook, event_type)
             if result is not None:
-                return result
-        return True
+                mode_tag_result: Final[bool] = bool(cast(object, result))  # cast-ok: helper return
+                return mode_tag_result and self.stream_scope_allows(cast(object, data), event_type)  # cast-ok: data
+        return self.stream_scope_allows(cast(object, data), event_type)  # cast-ok: data param
 
     def _event_hook_is_event_type(self, event_type: GuardrailEventHooks) -> bool:
         """

litellm/litellm_core_utils/realtime_streaming.py

@@ -782,7 +782,7 @@ def _has_realtime_guardrails_for_event_hooks(
             isinstance(cb, CustomGuardrail)
             and any(
                 cb.should_run_guardrail(
-                    data=self.request_data,
+                    data={**self.request_data, "stream": True},
                     event_type=et,
                 )
                 for et in event_hooks
@@ -847,7 +847,7 @@ async def run_realtime_guardrails(
         if event_hooks is None:
             event_hooks = [GuardrailEventHooks.realtime_input_transcription]
         _realtime_event_types: Final = event_hooks
-        _check_data: Final = {**self.request_data, "transcript": transcript}
+        _check_data: Final = {**self.request_data, "transcript": transcript, "stream": True}
         _already_run: Final[set] = set()
 
         for callback in litellm.callbacks:

litellm/llms/bedrock/passthrough/transformation.py

@@ -22,6 +22,13 @@
     from litellm.types.utils import CostResponseTypes
 
 
+BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"})
+
+
+def is_bedrock_streaming_endpoint(endpoint: str) -> bool:
+    return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS
+
+
 _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"})
 
 

litellm/llms/pass_through/guardrail_translation/handler.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, Any, Final, Optional
 
 from litellm._logging import verbose_proxy_logger
+from litellm.integrations.custom_guardrail import without_server_streaming_classification
 from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
 from litellm.proxy._types import PassThroughGuardrailSettings
 from litellm.types.utils import GenericGuardrailAPIInputs
@@ -80,7 +81,9 @@ def _extract_text_for_guardrail(
         from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
         payload_to_check: Final = {
-            k: v for k, v in data.items() if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj")
+            k: v
+            for k, v in without_server_streaming_classification(data).items()
+            if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj")
         }
         verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail")
         return safe_dumps(payload_to_check)

litellm/proxy/_lazy_openapi_snapshot.json

@@ -11858,6 +11858,34 @@
               "description": "When True (default), after sensitive data is detected and routed, all subsequent requests in the same session will continue routing to the same model.",
               "title": "Sticky Session Routing"
             },
+            "stream_scope": {
+              "anyOf": [
+                {
+                  "enum": [
+                    "streaming",
+                    "non_streaming",
+                    "both"
+                  ],
+                  "type": "string"
+                },
+                {
+                  "additionalProperties": {
+                    "enum": [
+                      "streaming",
+                      "non_streaming",
+                      "both"
+                    ],
+                    "type": "string"
+                  },
+                  "type": "object"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.",
+              "title": "Stream Scope"
+            },
             "template_id": {
               "anyOf": [
                 {
@@ -14783,6 +14811,34 @@
               "description": "When True (default), after sensitive data is detected and routed, all subsequent requests in the same session will continue routing to the same model.",
               "title": "Sticky Session Routing"
             },
+            "stream_scope": {
+              "anyOf": [
+                {
+                  "enum": [
+                    "streaming",
+                    "non_streaming",
+                    "both"
+                  ],
+                  "type": "string"
+                },
+                {
+                  "additionalProperties": {
+                    "enum": [
+                      "streaming",
+                      "non_streaming",
+                      "both"
+                    ],
+                    "type": "string"
+                  },
+                  "type": "object"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Whether this guardrail runs on streaming requests, non-streaming requests, or both. A string applies to every configured mode. A map overrides named modes (pre_call, during_call, post_call, ...); omitted keys default to both. Unset means both, matching historical behavior.",
+              "title": "Stream Scope"
+            },
             "template_id": {
               "anyOf": [
                 {

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -66,6 +66,7 @@
     PresidioPresidioConfigModelUserInterface,
     SupportedGuardrailIntegrations,
     ToolPermissionGuardrailConfigModel,
+    with_tolerated_stream_scope,
 )
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.proxy.guardrails.guardrail_hooks.hide_secrets import (
@@ -143,7 +144,7 @@ def _get_guardrails_list_response(
             GuardrailInfoResponse(
                 guardrail_id=guardrail.get("guardrail_id"),
                 guardrail_name=guardrail.get("guardrail_name"),
-                litellm_params=masked_params,
+                litellm_params=with_tolerated_stream_scope(masked_params),
                 guardrail_info=guardrail.get("guardrail_info"),
             )
         )
@@ -286,7 +287,7 @@ async def list_guardrails_v2(
             )
             masked_litellm_params = (
                 parse_tolerant_litellm_params(
-                    masked_litellm_params_dict,
+                    with_tolerated_stream_scope(masked_litellm_params_dict),
                     guardrail.get("guardrail_name") or "Unknown",
                     params_model=BaseLitellmParams,
                 )
@@ -333,7 +334,7 @@ async def list_guardrails_v2(
             )
             masked_in_memory_litellm_params_typed = (
                 parse_tolerant_litellm_params(
-                    masked_in_memory_litellm_params,
+                    with_tolerated_stream_scope(masked_in_memory_litellm_params),
                     guardrail.get("guardrail_name") or "Unknown",
                     params_model=BaseLitellmParams,
                 )
@@ -1260,7 +1261,7 @@ async def patch_guardrail(
         # Update litellm_params if default_on is provided or pii_entities_config is provided
         existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {})))
         current_litellm_params: Final = parse_tolerant_litellm_params(
-            existing_litellm_params,
+            with_tolerated_stream_scope(existing_litellm_params),
             existing_guardrail.get("guardrail_name") or "Unknown",
         )
         requested_litellm_params: Final[Mapping[str, object]] = (
@@ -1272,7 +1273,7 @@ async def patch_guardrail(
             MappingProxyType({**current_litellm_params.model_dump(exclude_unset=True), **requested_litellm_params})
         )
         try:
-            parsed_litellm_params: Final = LitellmParams(**merged_litellm_params)
+            parsed_litellm_params: Final = LitellmParams(**with_tolerated_stream_scope(merged_litellm_params))
         except ValidationError as validation_error:
             raise HTTPException(
                 status_code=422,
@@ -1433,7 +1434,7 @@ async def get_guardrail_info(guardrail_id: str):
         )
         masked_litellm_params = (
             parse_tolerant_litellm_params(
-                masked_litellm_params_dict,
+                with_tolerated_stream_scope(masked_litellm_params_dict),
                 result.get("guardrail_name") or "Unknown",
                 params_model=BaseLitellmParams,
             )