diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py index 1eea74a1e68..54138c51352 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py @@ -46,6 +46,12 @@ def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardra application_id=litellm_params.application_id, monitor_mode=litellm_params.monitor_mode, block_failures=litellm_params.block_failures, + streaming_end_of_stream_only=getattr( + litellm_params, "streaming_end_of_stream_only", None + ) + or False, + streaming_sampling_rate=getattr(litellm_params, "streaming_sampling_rate", None) + or 5, event_hook=litellm_params.mode, default_on=litellm_params.default_on, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 6aeaac949a9..9b26fd776ac 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -51,8 +51,12 @@ class NomaV2Guardrail(CustomGuardrail): application_id: Optional[str] = None, monitor_mode: Optional[bool] = None, block_failures: Optional[bool] = None, + streaming_end_of_stream_only: bool = False, + streaming_sampling_rate: int = 5, **kwargs: Any, ) -> None: + self.streaming_end_of_stream_only = streaming_end_of_stream_only + self.streaming_sampling_rate = streaming_sampling_rate self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py index c6fd587abe6..727c1d9500f 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py @@ -22,6 +22,14 @@ class NomaGuardrailConfigModel(GuardrailConfigModel): default=None, description="The Noma Application ID. Reads from NOMA_APPLICATION_ID env var if None.", ) + streaming_end_of_stream_only: Optional[bool] = Field( + default=None, + description="When true, only scan the final assembled stream chunk instead of intermediate chunks.", + ) + streaming_sampling_rate: Optional[int] = Field( + default=None, + description="Scan every Nth streaming chunk (default 5). Ignored when streaming_end_of_stream_only is true.", + ) @staticmethod def ui_friendly_name() -> str: @@ -49,6 +57,14 @@ class NomaV2GuardrailConfigModel(GuardrailConfigModel): default=None, description="When true, fail closed on Noma API errors.", ) + streaming_end_of_stream_only: Optional[bool] = Field( + default=None, + description="When true, only scan the final assembled stream chunk instead of intermediate chunks.", + ) + streaming_sampling_rate: Optional[int] = Field( + default=None, + description="Scan every Nth streaming chunk (default 5). Ignored when streaming_end_of_stream_only is true.", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma_v2.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma_v2.py index b6445a7c90d..2c8a6be2d29 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma_v2.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma_v2.py @@ -655,3 +655,68 @@ class TestNomaV2ApplicationIdResolution: payload = call_mock.call_args.kwargs["payload"] assert "application_id" not in payload + + +class TestNomaV2StreamingKnobs: + def test_streaming_knobs_default_values(self): + guardrail = NomaV2Guardrail( + api_key="test-api-key", + guardrail_name="test", + event_hook="pre_call", + default_on=True, + ) + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 5 + + def test_streaming_knobs_custom_values(self): + guardrail = NomaV2Guardrail( + api_key="test-api-key", + streaming_end_of_stream_only=True, + streaming_sampling_rate=10, + guardrail_name="test", + event_hook="pre_call", + default_on=True, + ) + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 10 + + def test_config_model_has_streaming_fields(self): + model = NomaV2GuardrailConfigModel( + streaming_end_of_stream_only=True, + streaming_sampling_rate=3, + ) + assert model.streaming_end_of_stream_only is True + assert model.streaming_sampling_rate == 3 + + def test_config_model_streaming_fields_default_to_none(self): + model = NomaV2GuardrailConfigModel() + assert model.streaming_end_of_stream_only is None + assert model.streaming_sampling_rate is None + + def test_initialize_guardrail_v2_passes_streaming_knobs(self): + from unittest.mock import patch as _patch + + from litellm.proxy.guardrails.guardrail_hooks.noma import ( + initialize_guardrail_v2, + ) + + class FakeLitellmParams: + api_key = "test-key" + api_base = "https://self-managed.local" + application_id = "app-1" + monitor_mode = False + block_failures = False + streaming_end_of_stream_only = True + streaming_sampling_rate = 8 + mode = "pre_call" + default_on = True + + guardrail = {"guardrail_name": "test-guardrail"} + + with _patch("litellm.logging_callback_manager.add_litellm_callback"): + result = initialize_guardrail_v2( + litellm_params=FakeLitellmParams(), guardrail=guardrail + ) + + assert result.streaming_end_of_stream_only is True + assert result.streaming_sampling_rate == 8