fix(noma): prevent falsy-coercion of streaming_sampling_rate=0

Replace `or 5`/`or False` with `is not None` check so that explicit
0 and False values are not silently replaced with defaults.
Add test for the rate=0 edge case and a comment clarifying that
streaming knobs are consumed by UnifiedLLMGuardrails framework.
This commit is contained in:
yryzhan 2026-05-20 18:31:58 +02:00
parent 6b2f44873f
commit 7953432e51
3 changed files with 39 additions and 6 deletions

View file

@ -46,12 +46,18 @@ 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,
streaming_end_of_stream_only=(
_v
if (_v := getattr(litellm_params, "streaming_end_of_stream_only", None))
is not None
else False
),
streaming_sampling_rate=(
_r
if (_r := getattr(litellm_params, "streaming_sampling_rate", None))
is not None
else 5
),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)

View file

@ -55,6 +55,7 @@ class NomaV2Guardrail(CustomGuardrail):
streaming_sampling_rate: int = 5,
**kwargs: Any,
) -> None:
# Read by UnifiedLLMGuardrails streaming hook
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(

View file

@ -693,6 +693,32 @@ class TestNomaV2StreamingKnobs:
assert model.streaming_end_of_stream_only is None
assert model.streaming_sampling_rate is None
def test_streaming_sampling_rate_zero_is_not_coerced(self):
"""streaming_sampling_rate=0 must not be silently replaced with default 5."""
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 = False
streaming_sampling_rate = 0
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_sampling_rate == 0
def test_initialize_guardrail_v2_passes_streaming_knobs(self):
from unittest.mock import patch as _patch