mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
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:
parent
6b2f44873f
commit
7953432e51
3 changed files with 39 additions and 6 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue