perf(responses): batch guardrail checks during streaming (#5664)

# What does this PR do?

Fixes a performance cliff when output guardrails are enabled on
streaming responses. Previously, every streaming token triggered an O(n)
string join, a `list_shields()` lookup, and a Safety API
`run_moderation()` call. For a 1000-token response this meant ~1000
redundant API calls and quadratic string reconstruction.

This PR:
- Extracts `resolve_guardrail_model_ids()` to cache shield lookups once
per request
- Batches guardrail checks every 200 characters (configurable via
`GUARDRAIL_BATCH_CHARS` env var) instead of every token
- Adds a final guardrail check at stream end for remaining buffered
content
- Flushes reasoning-only deltas per chunk so they stream in real time

## Test Plan

1. Unit tests pass (`236 passed`):
```bash
uv run pytest tests/unit/providers/responses/ -x --tb=short -q
```

2. New test verifies reasoning events stream without waiting for text
accumulation:
```bash
uv run pytest tests/unit/providers/inline/responses/builtin/responses/test_streaming.py::test_guardrailed_reasoning_streams_before_completion -v
```

3. Benchmark script for A/B testing against a running OGX server:
```bash
# Start server with per-token checking (before):
GUARDRAIL_BATCH_CHARS=1 SAFETY_MODEL=ollama/llama-guard3:1b uv run ogx stack run starter

# Run benchmark:
uv run python scripts/benchmark_guardrail_batching.py --model openai/gpt-4.1-nano

# Restart server with batched checking (after, default):
SAFETY_MODEL=ollama/llama-guard3:1b uv run ogx stack run starter

# Run benchmark again and compare
uv run python scripts/benchmark_guardrail_batching.py --model openai/gpt-4.1-nano
```

---------

Signed-off-by: Sébastien Han <seb@redhat.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: leseb

src/ogx/providers/inline/responses/builtin/responses/streaming.py

@@ -117,6 +117,7 @@
     convert_chat_choice_to_response_message,
     convert_mcp_tool_choice,
     is_function_tool_call,
+    resolve_guardrail_model_ids,
     run_guardrails,
     should_summarize_reasoning,
     summarize_reasoning,
@@ -129,6 +130,8 @@
 # Anything else is either a registered function tool (client-side) or a hallucinated name.
 _SERVER_SIDE_BUILTIN_TOOL_NAMES = frozenset({"web_search", "knowledge_search", "file_search"})
 
+_GUARDRAIL_BATCH_CHARS = 200
+
 # Maps OpenAI Chat Completions error codes to Responses API error codes
 _RESPONSES_API_ERROR_CODES = {
     "invalid_base64": "invalid_base64_image",
@@ -304,6 +307,7 @@ def __init__(
         self.accumulated_usage: OpenAIResponseUsage | None = None
         # Track if we've sent a refusal response
         self.violation_detected = False
+        self._guardrail_model_ids: list[str] = []
         # Track total calls made to built-in tools
         self.accumulated_builtin_tool_calls = 0
         # Track total output tokens generated across inference calls
@@ -411,8 +415,15 @@ async def create_response(self) -> AsyncIterator[OpenAIResponseObjectStream]:
 
         # Input safety validation - check messages before processing
         if self.guardrail_ids:
+            if self.safety_api is not None:
+                self._guardrail_model_ids = await resolve_guardrail_model_ids(self.safety_api, self.guardrail_ids)
             combined_text = interleaved_content_as_str([msg.content for msg in self.ctx.messages])
-            input_violation_message = await run_guardrails(self.safety_api, combined_text, self.guardrail_ids)
+            input_violation_message = await run_guardrails(
+                self.safety_api,
+                combined_text,
+                self.guardrail_ids,
+                model_ids=self._guardrail_model_ids,
+            )
             if input_violation_message:
                 logger.info("Input guardrail violation", input_violation_message=input_violation_message)
                 yield await self._create_refusal_response(input_violation_message)
@@ -1038,6 +1049,8 @@ async def _process_streaming_chunks(
         message_output_index = len(output_messages)
         reasoning_text_accumulated = []
         refusal_text_accumulated = []
+        pending_guardrail_events: list[OpenAIResponseObjectStream] = []
+        chars_since_last_check = 0
 
         async for raw_chunk in completion_result:
             # Providers returning OpenAIChatCompletionChunkWithReasoning wrap
@@ -1059,9 +1072,6 @@ async def _process_streaming_chunks(
             # Accumulate usage from chunks (typically in final chunk with stream_options)
             self._accumulate_chunk_usage(chunk)
 
-            # Track deltas for this specific chunk for guardrail validation
-            chunk_events: list[OpenAIResponseObjectStream] = []
-
             for chunk_choice in chunk.choices:
                 # Collect logprobs if present
                 chunk_logprobs = None
@@ -1115,12 +1125,14 @@ async def _process_streaming_chunks(
                     )
                     # Buffer text delta events for guardrail check
                     if self.guardrail_ids:
-                        chunk_events.append(text_delta_event)
+                        pending_guardrail_events.append(text_delta_event)
                     else:
                         yield text_delta_event
 
                 # Collect content for final response
-                chat_response_content.append(chunk_choice.delta.content or "")
+                content_delta = chunk_choice.delta.content or ""
+                chat_response_content.append(content_delta)
+                chars_since_last_check += len(content_delta)
                 if chunk_choice.finish_reason:
                     chunk_finish_reason = chunk_choice.finish_reason
 
@@ -1137,7 +1149,7 @@ async def _process_streaming_chunks(
                     ):
                         # Buffer reasoning events for guardrail check
                         if self.guardrail_ids:
-                            chunk_events.append(event)
+                            pending_guardrail_events.append(event)
                         else:
                             yield event
                     reasoning_part_emitted = True
@@ -1232,21 +1244,49 @@ async def _process_streaming_chunks(
                                     response_tool_call.function.arguments or ""
                                 ) + tool_call.function.arguments
 
-            # Output Safety Validation for this chunk
-            if self.guardrail_ids:
-                # Check guardrails on accumulated text so far
+            # Batched output safety validation. If we have only buffered reasoning events and
+            # no assistant text yet, flush per chunk so reasoning can stream in real time.
+            guardrail_check_due = chars_since_last_check >= _GUARDRAIL_BATCH_CHARS
+            if pending_guardrail_events and not any(chat_response_content):
+                guardrail_check_due = True
+
+            if self.guardrail_ids and guardrail_check_due:
                 accumulated_text = "".join(chat_response_content)
-                violation_message = await run_guardrails(self.safety_api, accumulated_text, self.guardrail_ids)
+                violation_message = await run_guardrails(
+                    self.safety_api,
+                    accumulated_text,
+                    self.guardrail_ids,
+                    model_ids=self._guardrail_model_ids,
+                )
                 if violation_message:
                     logger.info("Output guardrail violation", violation_message=violation_message)
-                    chunk_events.clear()
+                    pending_guardrail_events.clear()
                     yield await self._create_refusal_response(violation_message)
                     self.violation_detected = True
                     return
-                else:
-                    # No violation detected, emit all content events for this chunk
-                    for event in chunk_events:
-                        yield event
+                for event in pending_guardrail_events:
+                    yield event
+                pending_guardrail_events.clear()
+                chars_since_last_check = 0
+
+        # Final guardrail check on remaining buffered content
+        if self.guardrail_ids and pending_guardrail_events:
+            accumulated_text = "".join(chat_response_content)
+            violation_message = await run_guardrails(
+                self.safety_api,
+                accumulated_text,
+                self.guardrail_ids,
+                model_ids=self._guardrail_model_ids,
+            )
+            if violation_message:
+                logger.info("Output guardrail violation", violation_message=violation_message)
+                pending_guardrail_events.clear()
+                yield await self._create_refusal_response(violation_message)
+                self.violation_detected = True
+                return
+            for event in pending_guardrail_events:
+                yield event
+            pending_guardrail_events.clear()
 
         # Emit arguments.done events for completed tool calls (differentiate between MCP and function calls)
         for tool_call_index in sorted(chat_response_tool_calls.keys()):

src/ogx/providers/inline/responses/builtin/responses/utils.py

@@ -548,27 +548,38 @@ def is_function_tool_call(
     return False
 
 
-async def run_guardrails(safety_api: Safety | None, messages: str, guardrail_ids: list[str]) -> str | None:
-    """Run guardrails against messages and return violation message if blocked."""
-    if not messages:
-        return None
-
-    # If safety API is not available, skip guardrails
-    if safety_api is None:
-        return None
+async def resolve_guardrail_model_ids(safety_api: Safety, guardrail_ids: list[str]) -> list[str]:
+    """Resolve guardrail identifiers to concrete shield model IDs.
 
-    # Look up shields to get their provider_resource_id (actual model ID)
-    model_ids = []
+    Call once and pass the result to run_guardrails() to avoid repeated lookups.
+    """
     # TODO: list_shields not in Safety interface but available at runtime via API routing
     shields_list = await safety_api.routing_table.list_shields()  # type: ignore[attr-defined]
-
+    model_ids = []
     for guardrail_id in guardrail_ids:
         matching_shields = [shield for shield in shields_list.data if shield.identifier == guardrail_id]
         if matching_shields:
-            model_id = matching_shields[0].provider_resource_id
-            model_ids.append(model_id)
+            model_ids.append(matching_shields[0].provider_resource_id)
         else:
             raise ValueError(f"No shield found with identifier '{guardrail_id}'")
+    return model_ids
+
+
+async def run_guardrails(
+    safety_api: Safety | None,
+    messages: str,
+    guardrail_ids: list[str],
+    model_ids: list[str] | None = None,
+) -> str | None:
+    """Run guardrails against messages and return violation message if blocked."""
+    if not messages:
+        return None
+
+    if safety_api is None:
+        return None
+
+    if model_ids is None:
+        model_ids = await resolve_guardrail_model_ids(safety_api, guardrail_ids)
 
     guardrail_tasks = [
         safety_api.run_moderation(RunModerationRequest(input=messages, model=model_id)) for model_id in model_ids

tests/integration/agents/recordings/72c53cb1f81e5b2835ea301e84dbd3431d5f46ba974ee9cb5df3ff3f14d90732.json

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
 from collections.abc import AsyncIterator
 from unittest.mock import AsyncMock, MagicMock
 
@@ -23,11 +24,15 @@
 from ogx_api.inference.models import (
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAIChatCompletionChunkWithReasoning,
     OpenAIChatCompletionResponseMessage,
     OpenAIChatCompletionToolCall,
     OpenAIChatCompletionToolCallFunction,
     OpenAIChatCompletionUsage,
     OpenAIChoice,
+    OpenAIChoiceDelta,
+    OpenAIChunkChoice,
 )
 from ogx_api.openai_responses import (
     OpenAIResponseInputToolMCP,
@@ -577,3 +582,54 @@ async def test_uses_correct_summary_mode(self):
         call_args = mock_inference.openai_chat_completion.call_args[0][0]
         user_msg = call_args.messages[1].content
         assert "Preserve the key logical steps" in user_msg
+
+
+async def test_guardrailed_reasoning_streams_before_completion(mock_inference_api, mock_context, mock_safety_api):
+    """Guardrail batching should not buffer reasoning-only deltas until stream completion."""
+    mock_context.model = "test-model"
+    mock_context.temperature = None
+    mock_context.top_p = None
+    mock_context.frequency_penalty = None
+
+    orchestrator = StreamingResponseOrchestrator(
+        inference_api=mock_inference_api,
+        ctx=mock_context,
+        response_id="resp_reasoning_guardrails",
+        created_at=0,
+        text=MagicMock(),
+        max_infer_iters=1,
+        tool_executor=MagicMock(),
+        instructions=None,
+        safety_api=mock_safety_api,
+        guardrail_ids=["llama-guard"],
+    )
+
+    gate = asyncio.Event()
+
+    async def completion_result() -> AsyncIterator[OpenAIChatCompletionChunkWithReasoning]:
+        chunk = OpenAIChatCompletionChunk(
+            id="chatcmpl_reasoning",
+            choices=[
+                OpenAIChunkChoice(
+                    index=0,
+                    delta=OpenAIChoiceDelta(content=None, role="assistant"),
+                    finish_reason=None,
+                )
+            ],
+            created=1,
+            model="test-model",
+            object="chat.completion.chunk",
+        )
+        yield OpenAIChatCompletionChunkWithReasoning(chunk=chunk, reasoning_content="thinking...")
+
+        await gate.wait()
+
+    stream = orchestrator._process_streaming_chunks(completion_result(), output_messages=[])
+
+    # If reasoning is buffered until completion, this call will time out.
+    first_event = await asyncio.wait_for(anext(stream), timeout=0.5)
+    assert first_event.type in {"response.content_part.added", "response.reasoning_text.delta"}
+
+    gate.set()
+    async for _ in stream:
+        pass