fix(guardrails): forward mode and streaming params to crowdstrike_aidr handler (#39317)

* fix(guardrails): forward mode and streaming params to crowdstrike_aidr handler

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): drop stream_chunk_builder patch from crowdstrike cadence test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(guardrails): type test params and cover unsupported crowdstrike mode rejection

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -1,8 +1,8 @@
 from typing import TYPE_CHECKING, Final
 
-from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
+from litellm.types.guardrails import SupportedGuardrailIntegrations
 
-from .crowdstrike_aidr import CrowdStrikeAIDRHandler
+from .crowdstrike_aidr import CrowdStrikeAIDRHandler, streaming_params_from_litellm_params
 
 if TYPE_CHECKING:
     from litellm.types.guardrails import Guardrail, LitellmParams
@@ -15,17 +15,16 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
     if not guardrail_name:
         raise ValueError("CrowdStrike AIDR guardrail name is required")
 
+    streaming_params: Final = streaming_params_from_litellm_params(litellm_params)
     _crowdstrike_aidr_callback: Final = CrowdStrikeAIDRHandler(
         guardrail_name=guardrail_name,
         api_base=litellm_params.api_base,
         api_key=litellm_params.api_key,
-        # Exclude during_call to prevent duplicate input events
-        event_hook=[
-            GuardrailEventHooks.pre_call.value,
-            GuardrailEventHooks.post_call.value,
-        ],
+        event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
         fail_on_error=litellm_params.fail_on_error,
+        streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
+        streaming_sampling_rate=streaming_params.streaming_sampling_rate,
     )
     litellm.logging_callback_manager.add_litellm_callback(_crowdstrike_aidr_callback)
 

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py

@@ -24,8 +24,11 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.proxy.common_utils.callback_utils import (
     add_guardrail_to_applied_guardrails_header,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
 from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionToolParam
+from litellm.types.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import (
+    CrowdStrikeAIDRGuardrailConfigModelOptionalParams,
+)
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
@@ -153,6 +156,21 @@ def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Mapping[str, Any] |
     return merged if present else None
 
 
+def streaming_params_from_litellm_params(
+    litellm_params: LitellmParams,
+) -> CrowdStrikeAIDRGuardrailConfigModelOptionalParams:
+    extras: Final[Mapping[str, object]] = litellm_params.model_extra or {}
+    nested: Final = litellm_params.optional_params
+    optional_params: Final[Mapping[str, object]] = {} if nested is None else nested.model_dump()
+    return CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_validate(
+        {
+            name: value
+            for name in CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_fields
+            if (value := optional_params.get(name, extras.get(name))) is not None
+        }
+    )
+
+
 def _messages_since_last_assistant(
     messages: Sequence[AllMessageValues],
 ) -> _FilteredMessages:
@@ -241,6 +259,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         api_key: str | None = None,
         api_base: str | None = None,
         fail_on_error: bool | None = True,
+        streaming_end_of_stream_only: bool | None = None,
+        streaming_sampling_rate: int | None = None,
         **kwargs,
     ) -> None:
         """
@@ -250,10 +270,19 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
             guardrail_name (str): The name of the guardrail instance.
             api_key (str | None): The CrowdStrike AIDR API key. Reads from CS_AIDR_TOKEN env var if None.
             api_base (str | None): The CrowdStrike AIDR API base URL. Reads from CS_AIDR_BASE_URL env var if None.
+            streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of
+                every streaming_sampling_rate chunks. Defaults to False.
+            streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5.
             **kwargs: Additional arguments passed to the CustomGuardrail base class.
         """
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
         self.fail_on_error = True if fail_on_error is None else fail_on_error
+        self._set_streaming_params(
+            CrowdStrikeAIDRGuardrailConfigModelOptionalParams(
+                streaming_end_of_stream_only=streaming_end_of_stream_only,
+                streaming_sampling_rate=streaming_sampling_rate,
+            )
+        )
 
         self.api_key = api_key or os.environ.get("CS_AIDR_TOKEN")
         if not self.api_key:
@@ -274,6 +303,15 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
             "Initialized CrowdStrike AIDR Guardrail: name=%s, api_base=%s", guardrail_name, self.api_base
         )
 
+    def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None:
+        self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False
+        self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5
+
+    @override
+    def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
+        super().update_in_memory_litellm_params(litellm_params)
+        self._set_streaming_params(streaming_params_from_litellm_params(litellm_params))
+
     async def _call_crowdstrike_aidr_guard(
         self, payload: dict[str, Any], hook_name: str
     ) -> _GuardChatCompletionsResult:

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

@@ -4,7 +4,18 @@ from .base import GuardrailConfigModel
 
 
 class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel):
-    pass
+    streaming_end_of_stream_only: bool | None = Field(
+        default=None,
+        description="If False (default when unset), post_call scans the accumulated streamed response every "
+        "streaming_sampling_rate chunks and an in-flight block stops the stream. If True, the guard runs once "
+        "over the assembled response at end of stream, so flagged content may already have reached the client.",
+    )
+    streaming_sampling_rate: int | None = Field(
+        default=None,
+        ge=1,
+        description="When streaming_end_of_stream_only is False, scan the accumulated streamed response every Nth "
+        "chunk. Defaults to 5 when unset.",
+    )
 
 
 class CrowdStrikeAIDRGuardrailConfigModel(GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams]):

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py

@@ -3,7 +3,9 @@ from unittest.mock import patch
 import httpx
 import pytest
 from fastapi import HTTPException
+from pydantic import ValidationError
 
+import litellm
 from litellm.exceptions import Timeout
 from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
 from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail
@@ -12,8 +14,8 @@ from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr
     CrowdStrikeAIDRHandler,
 )
 from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
-from litellm.types.guardrails import Guardrail, LitellmParams
-from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse
+from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams
+from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream
 
 
 @pytest.fixture
@@ -1578,3 +1580,139 @@ async def test_unparseable_transformed_response_fails_closed_under_fail_open() -
 
     assert exc_info.value.status_code == 500
     assert "failing closed" in exc_info.value.detail["error"]
+
+
+def _initialize_from_config(**litellm_params_kwargs: object) -> CrowdStrikeAIDRHandler:
+    litellm_params = LitellmParams(
+        guardrail="crowdstrike_aidr",
+        api_key="pts_crowdstrike_tokenid",
+        api_base="https://api.crowdstrike.com/aidr/aiguard",
+        default_on=True,
+        **litellm_params_kwargs,
+    )
+    guardrail = Guardrail(guardrail_name="crowdstrike-aidr-guard", litellm_params=litellm_params)
+    return initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail)
+
+
+@pytest.mark.parametrize(
+    ("mode", "runs_pre_call", "runs_post_call"),
+    [("post_call", False, True), ("pre_call", True, False), (["pre_call", "post_call"], True, True)],
+)
+def test_initialize_guardrail_honors_configured_mode(
+    mode: str | list[str], runs_pre_call: bool, runs_post_call: bool
+) -> None:
+    handler = _initialize_from_config(mode=mode)
+
+    assert handler.should_run_guardrail({}, GuardrailEventHooks.pre_call) is runs_pre_call
+    assert handler.should_run_guardrail({}, GuardrailEventHooks.post_call) is runs_post_call
+
+
+def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_hooks() -> None:
+    with pytest.raises(ValueError, match="during_call is not in the supported event hooks"):
+        _initialize_from_config(mode="during_call")
+
+
+def test_initialize_guardrail_defaults_streaming_params() -> None:
+    handler = _initialize_from_config(mode="post_call")
+
+    assert handler.streaming_end_of_stream_only is False
+    assert handler.streaming_sampling_rate == 5
+
+
+@pytest.mark.parametrize(
+    "configured",
+    [
+        {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50},
+        {"optional_params": {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}},
+    ],
+)
+def test_initialize_guardrail_forwards_streaming_params(configured: dict[str, object]) -> None:
+    handler = _initialize_from_config(mode="post_call", **configured)
+
+    assert handler.streaming_end_of_stream_only is True
+    assert handler.streaming_sampling_rate == 50
+
+
+def test_initialize_guardrail_rejects_non_positive_sampling_rate() -> None:
+    with pytest.raises(ValidationError):
+        _initialize_from_config(mode="post_call", streaming_sampling_rate=0)
+
+
+def test_update_in_memory_litellm_params_reapplies_streaming_params() -> None:
+    handler = _initialize_from_config(mode="post_call")
+
+    handler.update_in_memory_litellm_params(
+        LitellmParams(
+            guardrail="crowdstrike_aidr",
+            mode="post_call",
+            streaming_end_of_stream_only=True,
+            streaming_sampling_rate=7,
+        )
+    )
+
+    assert handler.streaming_end_of_stream_only is True
+    assert handler.streaming_sampling_rate == 7
+
+
+def _stream_chunk(content: str, finish_reason: str | None) -> ModelResponseStream:
+    return ModelResponseStream(
+        model="gpt-4",
+        choices=[
+            litellm.StreamingChoices(
+                index=0, delta=Delta(role="assistant", content=content), finish_reason=finish_reason
+            )
+        ],
+    )
+
+
+async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: list[str]) -> int:
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
+
+    async def stream():
+        for i, content in enumerate(chunk_texts):
+            yield _stream_chunk(content, "stop" if i == len(chunk_texts) - 1 else None)
+
+    calls = 0
+
+    def _allow(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(
+            status_code=200, json={"result": {"blocked": False, "transformed": False}}, request=request
+        )
+
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": handler,
+        "metadata": {"guardrails": ["crowdstrike-aidr-guard"]},
+    }
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_allow)) as client:
+        await handler.async_handler.close()
+        handler.async_handler.client = client
+        async for _ in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"),
+            response=stream(),
+            request_data=request_data,
+        ):
+            pass
+    return calls
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("configured", "expected_calls"),
+    [
+        ({}, 3),
+        ({"streaming_sampling_rate": 2}, 6),
+        ({"streaming_end_of_stream_only": True}, 1),
+        ({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1),
+    ],
+)
+async def test_streaming_params_from_config_control_output_scan_cadence(
+    configured: dict[str, object], expected_calls: int
+) -> None:
+    """10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once."""
+    handler = _initialize_from_config(mode="post_call", **configured)
+
+    assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls