diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py index 678d611fdc..e1d9a7ce50 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py @@ -15,6 +15,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" if not guardrail_name: raise ValueError("OpenAI Moderation: guardrail_name is required") + optional_params = getattr(litellm_params, "optional_params", None) + openai_moderation_guardrail = OpenAIModerationGuardrail( guardrail_name=guardrail_name, **{ @@ -24,6 +26,12 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" "default_on": litellm_params.default_on, "event_hook": litellm_params.mode, "model": litellm_params.model, + "streaming_end_of_stream_only": _get_config_value( + litellm_params, optional_params, "streaming_end_of_stream_only" + ), + "streaming_sampling_rate": _get_config_value( + litellm_params, optional_params, "streaming_sampling_rate" + ), }, ) @@ -32,6 +40,14 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return openai_moderation_guardrail +def _get_config_value(litellm_params, optional_params, attribute_name): + 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) + + guardrail_initializer_registry = { SupportedGuardrailIntegrations.OPENAI_MODERATION.value: initialize_guardrail, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 4ddeac9a20..7e6f3dac00 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -57,6 +57,8 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): model: Optional[ Literal["omni-moderation-latest", "text-moderation-latest"] ] = None, + streaming_end_of_stream_only: Optional[bool] = None, + streaming_sampling_rate: Optional[int] = None, **kwargs, ): """Initialize OpenAI Moderation guardrail handler.""" @@ -85,6 +87,17 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): model or "omni-moderation-latest" ) + # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook + # via getattr(guardrail_to_apply, "streaming_*", default). + self.streaming_end_of_stream_only: bool = ( + False + if streaming_end_of_stream_only is None + else streaming_end_of_stream_only + ) + self.streaming_sampling_rate: int = ( + 5 if streaming_sampling_rate is None else streaming_sampling_rate + ) + if not self.api_key: raise ValueError( "OpenAI Moderation: api_key is required. Set OPENAI_API_KEY environment variable or pass it in configuration." diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py index 7d81cf9fe0..0fcc0f2309 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py @@ -29,6 +29,16 @@ class OpenAIModerationGuardrailConfigModel(BaseOpenAIModerationGuardrailConfigMo description="OpenAI API base URL. Defaults to 'https://api.openai.com/v1'.", ) + streaming_end_of_stream_only: Optional[bool] = Field( + default=False, + description="If False (default), moderation runs on sampled chunks during the stream at the cadence set by streaming_sampling_rate, and an in-flight violation stops further chunks from streaming. If True, moderation 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=5, + description="When streaming_end_of_stream_only is False, moderation runs every Nth streamed chunk. Ignored when streaming_end_of_stream_only is True.", + ) + @staticmethod def ui_friendly_name() -> str: return "OpenAI Moderation" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index bccfb4a1cb..16b5cbe858 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -820,3 +820,63 @@ def test_openai_moderation_process_error_metadata_none_edge_case(): # Internal key cleaned up assert "_openai_moderation_response" not in request_data["metadata"] + + +@pytest.mark.asyncio +async def test_openai_moderation_guardrail_streaming_defaults(): + """Defaults match the unified dispatcher: sampled in-stream, every 5th chunk.""" + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail(guardrail_name="test") + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 5 + + +@pytest.mark.asyncio +async def test_openai_moderation_guardrail_streaming_overrides(): + """Constructor-level overrides for the two streaming flags are stored on self.""" + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test", + streaming_end_of_stream_only=False, + streaming_sampling_rate=3, + ) + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 3 + + +@pytest.mark.asyncio +async def test_openai_moderation_initialize_guardrail_forwards_streaming_flags(): + """initialize_guardrail forwards streaming knobs from litellm_params (extra='allow').""" + import litellm + from litellm.proxy.guardrails.guardrail_hooks.openai import ( + initialize_guardrail as openai_initialize_guardrail, + ) + from litellm.types.guardrails import ( + Guardrail, + LitellmParams, + SupportedGuardrailIntegrations, + ) + + litellm.logging_callback_manager._reset_all_callbacks() + try: + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + litellm_params = LitellmParams( + guardrail=SupportedGuardrailIntegrations.OPENAI_MODERATION, + api_key="test-key", + model="omni-moderation-latest", + mode="post_call", + streaming_end_of_stream_only=False, + streaming_sampling_rate=2, + ) + guardrail = openai_initialize_guardrail( + litellm_params=litellm_params, + guardrail=Guardrail( + guardrail_name="test-openai-moderation", + litellm_params=litellm_params, + ), + ) + + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 2 + finally: + litellm.logging_callback_manager._reset_all_callbacks() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 461e0cebfc..0358ca998a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -67,10 +67,7 @@ async def test_openai_moderation_guardrail_streaming_latency(): request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": openai_guardrail, - "metadata": { - "guardrails": ["test-openai-moderation"], - "guardrail_config": {"streaming_sampling_rate": 1}, - }, # Check every chunk for test + "metadata": {"guardrails": ["test-openai-moderation"]}, } chunks_received = 0 @@ -161,10 +158,7 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): request_data = { "messages": [{"role": "user", "content": "generate hate"}], "guardrail_to_apply": openai_guardrail, - "metadata": { - "guardrails": ["test-openai-moderation"], - "guardrail_config": {"streaming_sampling_rate": 1}, - }, + "metadata": {"guardrails": ["test-openai-moderation"]}, } # Should raise HTTPException @@ -242,10 +236,7 @@ async def test_openai_moderation_streaming_end_of_stream_request_data_passthroug request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": openai_guardrail, - "metadata": { - "guardrails": ["test-openai-moderation"], - "guardrail_config": {"streaming_sampling_rate": 1}, - }, + "metadata": {"guardrails": ["test-openai-moderation"]}, } with ( @@ -284,3 +275,230 @@ async def test_openai_moderation_streaming_end_of_stream_request_data_passthroug guardrail_resp, dict ), f"Expected full moderation response dict, got {type(guardrail_resp)}: {guardrail_resp}" assert "results" in guardrail_resp + + +def _make_stream_chunk(content: str, finish_reason=None): + """Build a real ModelResponseStream so the handler's isinstance checks pass.""" + import litellm + from litellm.types.utils import Delta + + return ModelResponseStream( + model="gpt-4", + choices=[ + litellm.StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +@pytest.mark.asyncio +async def test_openai_moderation_streaming_default_uses_sampled_cadence(): + """Default config samples every 5th streamed chunk and runs a final aggregate + pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream + calls, plus 1 final = 3 total. + """ + import litellm + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = MagicMock() + mock_mod_response.results = [] + + async def mock_stream(): + chunks_data = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_model_response = ModelResponse( + id="mock-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content="ABCDEFGHIJ"), + finish_reason="stop", + ) + ], + ) + + with ( + patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ) as patched_make_request, + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, + } + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert patched_make_request.await_count == 3, ( + f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), " + f"got {patched_make_request.await_count}" + ) + + +@pytest.mark.asyncio +async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moderation_once(): + """Opt-in streaming_end_of_stream_only=True skips in-stream sampling and runs + moderation once on the assembled response at end of stream. + """ + import litellm + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + streaming_end_of_stream_only=True, + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = MagicMock() + mock_mod_response.results = [] + + async def mock_stream(): + chunks_data = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_model_response = ModelResponse( + id="mock-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content="ABCDEFGHIJ"), + finish_reason="stop", + ) + ], + ) + + with ( + patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ) as patched_make_request, + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, + } + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert patched_make_request.await_count == 1, ( + f"Expected exactly one moderation call at end of stream, " + f"got {patched_make_request.await_count}" + ) + + +@pytest.mark.asyncio +async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled(): + """With streaming_end_of_stream_only=False and streaming_sampling_rate=2, + moderation runs every 2nd chunk during the stream, plus once more at end. + """ + import litellm + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + streaming_end_of_stream_only=False, + streaming_sampling_rate=2, + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = MagicMock() + mock_mod_response.results = [] + + async def mock_stream(): + chunks_data = ["A", "B", "C", "D", "E", "F"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_model_response = ModelResponse( + id="mock-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content="ABCDEF"), + finish_reason="stop", + ) + ], + ) + + with ( + patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ) as patched_make_request, + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, + } + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls), + # plus the final aggregate pass after the stream ends (1 call) = 4 total. + assert patched_make_request.await_count == 4, ( + f"Expected 4 moderation calls (3 sampled + 1 final aggregate), " + f"got {patched_make_request.await_count}" + )