fix(openai-moderation): wire streaming flags through to unified dispatcher (#27324)

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,
 }

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."

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"

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()

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}"
+        )