diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index 5e75b7d4d94..c88e6e97a96 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index c8284fac440..f7f500b1adc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/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: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index f47c38af3e3..6beca030a3a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/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]): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index ec7854b9a35..1b50ea53db2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/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