From 05711d3fa02651eef3225ab7efd08af3f33f949d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:18:21 +0000 Subject: [PATCH] fix(guardrails): honor streaming knobs in noma_v2 guardrail Wire streaming_end_of_stream_only and streaming_sampling_rate from config into NomaV2Guardrail so the unified streaming dispatcher can resolve them off the instance. --- .../guardrail_hooks/noma/__init__.py | 24 +++ .../guardrail_hooks/noma/noma_v2.py | 9 + .../proxy/guardrails/guardrail_hooks/noma.py | 17 ++ .../guardrail_hooks/test_noma_v2.py | 160 ++++++++++++++++++ 4 files changed, 210 insertions(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py index f0e6c6677a3..c10e71e1982 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py @@ -1,5 +1,7 @@ from typing import TYPE_CHECKING +from pydantic import TypeAdapter + from litellm.types.guardrails import SupportedGuardrailIntegrations from .noma import NomaGuardrail @@ -34,9 +36,29 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return _noma_callback +_END_OF_STREAM_ONLY_ADAPTER: TypeAdapter[bool | None] = TypeAdapter(bool | None) +_SAMPLING_RATE_ADAPTER: TypeAdapter[int | None] = TypeAdapter(int | None) + + +def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> object: + if optional_params is not None: + value = getattr(optional_params, attribute_name, None) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm + optional_params = getattr(litellm_params, "optional_params", None) + end_of_stream_only = _END_OF_STREAM_ONLY_ADAPTER.validate_python( + _get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only") + ) + sampling_rate = _SAMPLING_RATE_ADAPTER.validate_python( + _get_config_value(litellm_params, optional_params, "streaming_sampling_rate") + ) + _noma_v2_callback = NomaV2Guardrail( guardrail_name=guardrail.get("guardrail_name", ""), api_key=litellm_params.api_key, @@ -44,6 +66,8 @@ 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=end_of_stream_only, + streaming_sampling_rate=sampling_rate, 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 9cf0986c122..8bf1e11536b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -51,6 +51,8 @@ class NomaV2Guardrail(CustomGuardrail): application_id: Optional[str] = None, monitor_mode: Optional[bool] = None, block_failures: Optional[bool] = None, + streaming_end_of_stream_only: bool | None = None, + streaming_sampling_rate: int | None = None, **kwargs: Any, ) -> None: self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -68,6 +70,13 @@ class NomaV2Guardrail(CustomGuardrail): else: self.block_failures = block_failures + 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 + if self._requires_api_key(api_base=self.api_base) and not self.api_key: raise ValueError("Noma v2 guardrail requires api_key when using Noma SaaS endpoint") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py index c6fd587abe6..8461501a7f7 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py @@ -49,6 +49,23 @@ 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=( + "If False (default when unset), streaming post-call scans run on sampled chunks at the cadence set " + "by streaming_sampling_rate, and an in-flight block stops further chunks from streaming. If True, the " + "scan 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." + ), + ) + streaming_sampling_rate: Optional[int] = Field( + default=None, + ge=1, + description=( + "When streaming_end_of_stream_only is False, the streaming post-call scan runs every Nth streamed " + "chunk. Ignored when streaming_end_of_stream_only is True. Must be >= 1 when set; defaults to 5." + ), + ) @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..818a1516f60 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 @@ -3,6 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.noma import NomaV2Guardrail from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage from litellm.types.proxy.guardrails.guardrail_hooks.noma import ( @@ -655,3 +656,162 @@ class TestNomaV2ApplicationIdResolution: payload = call_mock.call_args.kwargs["payload"] assert "application_id" not in payload + + +class TestNomaV2StreamingKnobs: + @staticmethod + def _guardrail(**kwargs): + return NomaV2Guardrail( + api_key="test-api-key", + api_base="https://api.test.noma.security/", + application_id="test-app", + guardrail_name="test-noma-v2-guardrail", + event_hook="post_call", + default_on=True, + **kwargs, + ) + + @staticmethod + def _stream(num_chunks: int): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + async def gen(): + for i in range(num_chunks): + yield ModelResponseStream( + model="gpt-4o", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=f"chunk-{i}", role="assistant"), + finish_reason="stop" if i == num_chunks - 1 else None, + ) + ], + ) + + return gen() + + async def _run_stream(self, guardrail, num_chunks: int): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + scan_mock = AsyncMock(return_value={"action": "NONE"}) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + + with patch.object(guardrail, "_call_noma_scan", scan_mock): + chunks = [ + chunk + async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=self._stream(num_chunks), + request_data=request_data, + ) + ] + + return chunks, scan_mock.call_count + + def test_init_defaults_streaming_knobs(self): + guardrail = self._guardrail() + + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 5 + + def test_init_rejects_sampling_rate_below_one(self): + with pytest.raises(ValueError, match="streaming_sampling_rate must be >= 1"): + self._guardrail(streaming_sampling_rate=0) + + @pytest.mark.asyncio + async def test_streaming_end_of_stream_only_scans_once(self): + guardrail = self._guardrail(streaming_end_of_stream_only=True, streaming_sampling_rate=1) + + chunks, scan_calls = await self._run_stream(guardrail, num_chunks=6) + + assert len(chunks) == 6 + assert scan_calls == 1 + + @pytest.mark.asyncio + async def test_streaming_sampling_rate_controls_scan_cadence(self): + sparse_calls = (await self._run_stream(self._guardrail(streaming_sampling_rate=6), num_chunks=6))[1] + dense_calls = (await self._run_stream(self._guardrail(streaming_sampling_rate=2), num_chunks=6))[1] + + assert sparse_calls == 2 + assert dense_calls == 4 + + @pytest.mark.asyncio + async def test_default_streaming_scans_every_fifth_chunk(self): + guardrail = self._guardrail() + + chunks, scan_calls = await self._run_stream(guardrail, num_chunks=6) + + assert len(chunks) == 6 + assert scan_calls == 2 + + +class TestNomaV2StreamingKnobInitialization: + @staticmethod + def _initialize(guardrail_provider: str = "noma_v2", **litellm_param_overrides): + from litellm.proxy.guardrails.guardrail_hooks.noma import ( + guardrail_initializer_registry, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail=guardrail_provider, + mode="post_call", + api_key="test-api-key", + api_base="https://api.test.noma.security/", + **litellm_param_overrides, + ) + initializer = guardrail_initializer_registry[guardrail_provider] + return initializer(litellm_params=litellm_params, guardrail={"guardrail_name": "noma-v2"}) + + def test_forwards_streaming_knobs_from_litellm_params(self): + guardrail = self._initialize(streaming_end_of_stream_only=True, streaming_sampling_rate=3) + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 3 + + def test_forwards_streaming_knobs_from_string_values(self): + guardrail = self._initialize(streaming_end_of_stream_only="true", streaming_sampling_rate="3") + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 3 + + def test_forwards_streaming_knobs_via_legacy_noma_with_use_v2(self): + guardrail = self._initialize( + guardrail_provider="noma", + use_v2=True, + streaming_end_of_stream_only=True, + streaming_sampling_rate=4, + ) + + assert isinstance(guardrail, NomaV2Guardrail) + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 4 + + def test_forwards_streaming_knobs_from_optional_params(self): + guardrail = self._initialize(optional_params={"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}) + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 2 + + def test_defaults_when_streaming_knobs_absent(self): + guardrail = self._initialize() + + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 5 + + def test_rejects_invalid_sampling_rate_from_config(self): + with pytest.raises(ValueError): + self._initialize(streaming_sampling_rate=0) + + def test_config_model_declares_streaming_knobs(self): + fields = NomaV2GuardrailConfigModel.model_fields + + assert "streaming_end_of_stream_only" in fields + assert "streaming_sampling_rate" in fields