From 60d50d2ed7392b3d52080f0d7669a83bb20f63b5 Mon Sep 17 00:00:00 2001 From: Oren Hazai Date: Thu, 2 Jul 2026 18:29:42 +0300 Subject: [PATCH 1/5] Add streaming speculative generation (SG2) Run input rails concurrently with the LLM and output rails for stream_async, holding output-rail-validated chunks in a bounded awaiting-release buffer until the input verdict is known (check-first only). Input reject aborts generation; output-rail reject during speculation short-circuits and cancels the input task. - config: add rails.input.speculative_max_buffered_tokens (release-buffer bound with backpressure); document streaming support and the stream_first override - iorails: decouple input rails from the generation task into a concurrent _check_speculative_input_safety; add _gate_on_input release gate; force check-first for speculative streaming; preserve tool_input_rails vs input_rails violation param parity with the non-streaming path - telemetry: emit the full speculative signal set (durations, overlap, time_saved, cancellation_event, and streaming release-queue/output-rail signals) - tests: streaming pass/reject/output-early-reject/backpressure/no-output-rails/ tool-result-param coverage; replace the obsolete warn-and-fallback test - docs: streaming speculative-generation section and span-reference attributes --- .../yaml-schema/guardrails-configuration.mdx | 25 +- docs/observability/tracing/span-reference.mdx | 20 +- nemoguardrails/guardrails/iorails.py | 400 ++++++++++++++++-- nemoguardrails/guardrails/telemetry.py | 51 ++- nemoguardrails/rails/llm/config.py | 20 +- nemoguardrails/tracing/constants.py | 20 + tests/guardrails/test_data.py | 30 ++ tests/guardrails/test_iorails_streaming.py | 39 +- .../guardrails/test_speculative_generation.py | 255 ++++++++++- 9 files changed, 796 insertions(+), 64 deletions(-) diff --git a/docs/configure-rails/yaml-schema/guardrails-configuration.mdx b/docs/configure-rails/yaml-schema/guardrails-configuration.mdx index 42dbd562b9..696029022e 100644 --- a/docs/configure-rails/yaml-schema/guardrails-configuration.mdx +++ b/docs/configure-rails/yaml-schema/guardrails-configuration.mdx @@ -331,9 +331,12 @@ To decide whether Speculative Generation makes sense for your use-case, explore Speculative generation currently requires the opt-in IORails engine. To enable IORails, set `NEMO_GUARDRAILS_IORAILS_ENGINE=1`. Speculative - generation is supported only for non-streaming requests (`generate_async`). - When speculative generation is enabled, streaming requests (`stream_async`) - fall back to sequential execution and emit a warning. + generation is supported for both non-streaming requests (`generate_async`) + and streaming requests (`stream_async`). Streaming speculation only supports + check-first behavior: if `rails.output.streaming.stream_first` is `True`, it + is overridden to check-first for speculative requests (with a warning), since + tokens cannot reach the client before the input verdict is known. Streaming + speculation with output rails requires `rails.output.streaming.enabled: True`. ### How Speculative Generation Works @@ -366,6 +369,22 @@ The engine handles three outcomes: Output rails always run after the main LLM completes. Speculative generation does not change the output-rail path. +### Streaming Speculative Generation + +For streaming requests (`stream_async`), speculative generation runs the input rails, the main LLM, and the output rails concurrently. +While the input-rail verdict is still pending (the *speculation window*), validated tokens are held in a bounded in-memory buffer instead of being sent to the client — no token reaches the caller before the input is confirmed safe. + +1. Start the input rails and the main LLM stream in parallel. Output rails process tokens as they stream (check-first), at the LLM's natural generation rate. +2. Hold output-rail-validated tokens in the awaiting-release buffer. +3. Resolve the race: + - **Input rails pass:** flush the held buffer to the client and continue streaming normally. + - **Input rails reject:** discard the held buffer, abort the generation, and return the refusal message. No tokens are leaked. + - **Output rails reject during the speculation window:** return the refusal early and cancel the still-running input rails — the request is refused before the input verdict is known. + +Streaming speculation only supports check-first behavior. If `rails.output.streaming.stream_first` is `True`, it is overridden to check-first for speculative requests and a warning is emitted. + +The awaiting-release buffer is bounded by `rails.input.speculative_max_buffered_tokens` (default `4096`, counted in chunks). The buffer fills when the input rails are slower than the main LLM plus output rails, so validated chunks accumulate while the input verdict is still pending. When the bound is reached, the engine pauses speculative processing and waits for the input verdict before buffering more validated chunks (releasing on pass, tearing down on reject). This caps the release buffer only — the main LLM continues generating into the internal stream buffer — so total speculative memory is bounded by the model's finite output and the input-rail latency, not by stopping the backend. + ### Configuration Example To enable speculative generation, set `speculative_generation: True` under `rails.input` in the `config.yml` file. diff --git a/docs/observability/tracing/span-reference.mdx b/docs/observability/tracing/span-reference.mdx index ceec72ef1b..e2a2a27282 100644 --- a/docs/observability/tracing/span-reference.mdx +++ b/docs/observability/tracing/span-reference.mdx @@ -219,11 +219,21 @@ Every span sets ERROR status and records the exception when one propagates throu When speculative generation is active, the request span also carries these attributes: -| Attribute | Type | When set | Description | -| ---------------------------------------- | ------- | ------------------------- | --------------------------------------------------------------- | -| `speculative_generation.mode_active` | boolean | Speculative generation on | The value `true`. | -| `speculative_generation.first_completed` | string | Speculative generation on | The branch that completed first: `input_rails` or `generation`. | -| `speculative_generation.first_rejector` | string | Speculative generation on | The branch that rejected the request: `input_rails` or `none`. | +| Attribute | Type | When set | Description | +| ------------------------------------------------------ | ------- | ------------------------- | --------------------------------------------------------------------------------------------------------------------- | +| `speculative_generation.mode_active` | boolean | Speculative generation on | The value `true`. | +| `speculative_generation.first_completed` | string | Speculative generation on | The branch that completed first: `input_rails`, `generation`, or `output_rails`. | +| `speculative_generation.first_rejector` | string | Speculative generation on | The branch that rejected the request: `input_rails`, `output_rails`, or `none`. | +| `speculative_generation.rails_duration_ms` | float | Speculative generation on | Wall-clock time for the input rails. | +| `speculative_generation.generation_duration_ms` | float | Speculative generation on | Wall-clock time for the main LLM generation. | +| `speculative_generation.overlap_ms` | float | Speculative generation on | Time the input rails and generation ran concurrently. | +| `speculative_generation.time_saved_ms` | float | Speculative generation on | Estimated saving versus sequential execution (the overlap for safe requests, `0` for rejected requests). | +| `speculative_generation.cancellation_event` | string | Speculative generation on | Which task was cancelled: `none`, `generation_cancelled`, `input_rails_cancelled`, or `output_rails_cancelled`. | +| `speculative_generation.output_rails_early_reject` | boolean | Streaming speculation | Output rails rejected during the speculation window (short-circuit before the input verdict). | +| `speculative_generation.output_rails_speculation_chunks` | int | Streaming speculation | Number of chunks output rails processed during the speculation window. | +| `speculative_generation.output_rails_wasted_chunks` | int | Streaming speculation | Validated chunks discarded because the request was rejected. | +| `speculative_generation.release_queue_duration_ms` | float | Streaming speculation | Time validated tokens were held in the awaiting-release buffer. | +| `speculative_generation.release_queue_token_count` | int | Streaming speculation | Number of validated tokens held awaiting the input-rail verdict. | **`guardrails.rail`** is one span per rail that runs, wrapping the rail's execution. Span kind `INTERNAL`. diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 1b03d3d818..c3fc74d61f 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -27,7 +27,7 @@ import warnings from collections.abc import AsyncGenerator, AsyncIterator from contextlib import nullcontext, suppress -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union from nemoguardrails.actions.llm.utils import _extract_and_remove_think_tags from nemoguardrails.base_guardrails import BaseGuardrails @@ -324,6 +324,7 @@ def __init__(self, config: RailsConfig, *, _report_usage: bool = True) -> None: content_capture_enabled=self._content_capture_enabled, ) self._speculative_generation = config.rails.input.speculative_generation or False + self._speculative_max_buffered_tokens = config.rails.input.speculative_max_buffered_tokens # Non-streaming admission queue + worker pool (owned by IORails so # all request-path concurrency controls sit under one roof). The @@ -751,12 +752,25 @@ def stream_async( asyncio.QueueFull: If the streaming concurrency limit is reached (load shedding). """ - if self._speculative_generation: + self._validate_streaming_with_output_rails() + + # Speculative streaming (SG2): input rails race the LLM instead of blocking + # before it. Only check-first is supported — during the speculation window + # tokens cannot reach the client, so stream_first is overridden to + # check-first for speculative requests. + # + # NOTE: this precondition warning fires per stream_async() call (i.e. at + # request time), not once at engine startup. Flagged here so developers + # know the check-first override is decided per request, not globally. + use_speculative = self._speculative_generation + force_check_first = False + if use_speculative and self._has_streaming_output_rails and self.config.rails.output.streaming.stream_first: warnings.warn( - "speculative_generation is not supported for streaming; falling back to sequential", + "speculative_generation with stream_first=True is not supported for streaming; " + "forcing check-first behavior for this request", stacklevel=2, ) - self._validate_streaming_with_output_rails() + force_check_first = True if include_metadata and self._has_streaming_output_rails: raise ValueError( @@ -782,7 +796,7 @@ def stream_async( # (see ModelEngine.stream_call), so a plain rebind is sufficient. accumulated_tool_calls: list[ToolCall] = [] - async def _generation_task(request_span): + async def _generation_task(request_span, *, run_input_rails: bool = True, spec_stats: Optional[dict] = None): """Background task: input rails → stream LLM chunks → push to handler. ``request_span`` is the IORails request span (or ``None`` when @@ -791,43 +805,50 @@ async def _generation_task(request_span): ``trace.get_current_span()`` which could return the host app's ambient span and pollute unrelated traces. + When ``run_input_rails`` is False (speculative streaming), the + tool-result and input rails are skipped here — the caller runs them + in a concurrent task so LLM tokens start flowing immediately. When + ``spec_stats`` is provided, the task records its wall-clock duration + into it (``generation_duration_ms``) for speculative telemetry. + Inherits the request ID from the caller context via create_task(). """ nonlocal accumulated_tool_calls req_id = get_request_id() t0 = time.monotonic() try: - # Step 0: Tool-result rails. Client/agent harness executes tool calls and sends - # results of execution to Main LLM along with prior conversation history - # Symmetric with INPUT rails for dialog use-case - log.info("[%s] Running tool result rails", req_id) - tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) - if not tool_result.is_safe: - log.info("[%s] Tool result blocked: %s", req_id, tool_result.reason) - if self._metrics_enabled: - record_request_blocked(RailDirection.INPUT) - await streaming_handler.push_chunk( - self._guardrails_violation_payload( - f"Blocked by tool input rails: {tool_result.reason}", "tool_input_rails" + if run_input_rails: + # Step 0: Tool-result rails. Client/agent harness executes tool calls and sends + # results of execution to Main LLM along with prior conversation history + # Symmetric with INPUT rails for dialog use-case + log.info("[%s] Running tool result rails", req_id) + tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) + if not tool_result.is_safe: + log.info("[%s] Tool result blocked: %s", req_id, tool_result.reason) + if self._metrics_enabled: + record_request_blocked(RailDirection.INPUT) + await streaming_handler.push_chunk( + self._guardrails_violation_payload( + f"Blocked by tool input rails: {tool_result.reason}", "tool_input_rails" + ) ) - ) - await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore[arg-type] - return + await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore[arg-type] + return - # Step 1: Input rails (non-streaming) - log.info("[%s] Running input rails", req_id) - input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) - if not input_result.is_safe: - log.info("[%s] Input blocked: %s", req_id, input_result.reason) - if self._metrics_enabled: - record_request_blocked(RailDirection.INPUT) - await streaming_handler.push_chunk( - self._guardrails_violation_payload( - f"Blocked by input rails: {input_result.reason}", "input_rails" + # Step 1: Input rails (non-streaming) + log.info("[%s] Running input rails", req_id) + input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) + if not input_result.is_safe: + log.info("[%s] Input blocked: %s", req_id, input_result.reason) + if self._metrics_enabled: + record_request_blocked(RailDirection.INPUT) + await streaming_handler.push_chunk( + self._guardrails_violation_payload( + f"Blocked by input rails: {input_result.reason}", "input_rails" + ) ) - ) - await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore[arg-type] - return + await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore[arg-type] + return # Step 2: Stream main LLM content from structured response. # delta_content is forwarded as text chunks; delta_tool_calls are @@ -885,6 +906,8 @@ async def _generation_task(request_span): await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore[arg-type] finally: elapsed_ms = (time.monotonic() - t0) * 1000 + if spec_stats is not None: + spec_stats["generation_duration_ms"] = elapsed_ms log.info("[%s] generation task completed time=%.1fms", req_id, elapsed_ms) async def _wrapped_iterator(): @@ -939,22 +962,68 @@ async def _wrapped_iterator(): # the consumer, so the terminal tool-call chunk is # suppressed (never surface tool calls after a failure/block). error_emitted = False + # Speculative-streaming state, declared before the try so + # the outer finally can always reference them (defined even + # if the try body raises before assignment). + spec_stats: Optional[dict[str, Any]] = None + input_task: Optional[asyncio.Task] = None try: log.info("[%s] stream_async called", req_id) log.debug("[%s] stream_async messages=%s", req_id, truncate(messages)) - task = asyncio.create_task(_generation_task(request_span)) + # Speculative streaming (SG2): run input rails in a + # concurrent task so the LLM starts streaming right + # away, and skip input rails inside the generation + # task. spec_stats collects telemetry filled by the + # input task, the generation task, and the gate. + if use_speculative: + spec_stats = { + "first_completed": None, + "first_rejector": GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE, + "safe": True, + "output_rails_early_reject": False, + "cancellation_event": GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE, + } + input_task = asyncio.create_task( + self._check_speculative_input_safety( + messages, + input_enabled=input_enabled, + tool_input_enabled=tool_input_enabled, + spec_stats=spec_stats, + ) + ) + task = asyncio.create_task( + _generation_task(request_span, run_input_rails=False, spec_stats=spec_stats) + ) + else: + task = asyncio.create_task(_generation_task(request_span)) try: - # Determine base iterator: with or without output rails + # Determine the inner iterator: with or without output rails. if self._has_streaming_output_rails: - base_iterator = self._run_output_rails_in_streaming( + inner_iterator = self._run_output_rails_in_streaming( streaming_handler=streaming_handler, messages=messages, enabled=output_enabled, include_metadata=include_metadata, + force_check_first=force_check_first, ) else: - base_iterator = streaming_handler + # SG2 buffer-and-release: with no output rails configured, + # speculation still runs. Raw LLM tokens flow straight from + # the streaming handler; when speculating, the gate below holds + # them in the bounded release buffer until input rails pass, + # then flushes. + inner_iterator = streaming_handler + + # Gate raw/validated chunks on the input rails verdict during the + # speculation window; pass through unchanged when not speculating. + if use_speculative: + assert input_task is not None and spec_stats is not None + base_iterator = self._gate_on_input( + inner_iterator, input_task, spec_stats, include_metadata=include_metadata + ) + else: + base_iterator = inner_iterator async for chunk in base_iterator: if chunk is not None: @@ -1009,6 +1078,14 @@ async def _wrapped_iterator(): task.cancel() with suppress(asyncio.CancelledError): await task + # Defensive: the gate cancels+drains input_task in + # its own finally, but ensure it never leaks if the + # gate was never fully iterated (early break/error). + if input_task is not None: + if not input_task.done(): + input_task.cancel() + with suppress(asyncio.CancelledError, Exception): + await input_task except Exception: elapsed_ms = (time.monotonic() - t0) * 1000 log.error("[%s] stream_async failed time=%.1fms", req_id, elapsed_ms, exc_info=True) @@ -1025,11 +1102,251 @@ async def _wrapped_iterator(): if self._content_capture_enabled: output_text = "".join(delivered) if delivered else None set_request_content(request_span, messages, output_text) + # Stamp speculative-generation telemetry on the request + # span. Runs after teardown so both task durations are + # recorded (the generation task's finally has run once it + # was awaited above). overlap ≈ min(both durations) since + # both tasks start together; time_saved is the overlap for + # safe requests and 0 for rejected ones. + if use_speculative and spec_stats is not None: + rails_ms = spec_stats.get("rails_duration_ms") + gen_ms = spec_stats.get("generation_duration_ms") + overlap_ms = ( + min(rails_ms, gen_ms) if rails_ms is not None and gen_ms is not None else None + ) + time_saved_ms = ( + None if overlap_ms is None else (overlap_ms if spec_stats.get("safe") else 0.0) + ) + set_speculative_span_attrs( + request_span, + spec_stats.get("first_completed") + or GuardrailsAttributes.SPECULATIVE_FIRST_COMPLETED_GENERATION, + spec_stats.get( + "first_rejector", GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE + ), + rails_duration_ms=rails_ms, + generation_duration_ms=gen_ms, + overlap_ms=overlap_ms, + time_saved_ms=time_saved_ms, + cancellation_event=spec_stats.get("cancellation_event"), + output_rails_early_reject=spec_stats.get("output_rails_early_reject"), + output_rails_speculation_chunks=spec_stats.get( + "output_rails_speculation_chunks" + ), + output_rails_wasted_chunks=spec_stats.get("output_rails_wasted_chunks"), + release_queue_duration_ms=spec_stats.get("release_queue_duration_ms"), + release_queue_token_count=spec_stats.get("release_queue_token_count"), + ) finally: self._stream_semaphore.release() return _wrapped_iterator() + async def _check_speculative_input_safety( + self, + messages: LLMMessages, + *, + input_enabled: Union[bool, list[str]] = True, + tool_input_enabled: Union[bool, list[str]] = True, + spec_stats: Optional[dict] = None, + ): + """Concurrent input-safety check for speculative streaming (SG2). + + Runs tool-result rails then input rails, returning ``(RailResult, param)`` + for the first failing rail (or the safe input result). ``param`` + identifies the rail family for the client-facing violation payload + (``tool_input_rails`` vs ``input_rails``), matching the non-speculative + path. Runs as its own task so LLM generation can stream concurrently + during the speculation window. Records its wall-clock duration into + ``spec_stats['rails_duration_ms']`` for telemetry. + """ + t0 = time.monotonic() + try: + tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) + if not tool_result.is_safe: + return tool_result, "tool_input_rails" + input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) + return input_result, "input_rails" + finally: + if spec_stats is not None: + spec_stats["rails_duration_ms"] = (time.monotonic() - t0) * 1000 + + async def _gate_on_input( + self, + base_iterator: AsyncIterator[Union[str, dict]], + input_task: "asyncio.Task", + spec_stats: dict, + *, + include_metadata: Optional[bool] = False, + ) -> AsyncGenerator[Union[str, dict], None]: + """Hold streamed chunks until the input rails verdict is known (SG2). + + During the speculation window the LLM (and output rails, when + configured) run concurrently with input rails. Chunks that arrive + before the input verdict are held in a bounded in-memory buffer rather + than sent to the client — the input safety verdict is not yet known, so + nothing may reach the caller. On input PASS the held buffer is flushed + and streaming continues normally. On input REJECT (or an output-rails + early reject / generation error surfaced as an error chunk) the held + buffer is discarded and a refusal / the error payload is emitted. + + Cancellation (SG2): on input reject the generation task is torn down by + the caller's ``finally``; on an output-rails early reject the still- + running input rails task is cancelled here so the request aborts before + the input verdict arrives. + + When the held buffer reaches ``speculative_max_buffered_tokens`` the gate + stops consuming the base iterator and blocks on the input verdict, forcing + an early resolve (release on pass, teardown on reject). This bounds only + the release buffer — the background generation task keeps producing into + the stream queue — so total memory is bounded by the input-rail latency and + the model's finite output, not by halting generation. (An alternative + overflow policy, aborting the request when the bound is exceeded, would cap + total memory by cancelling the producer; backpressure is used here instead.) + """ + req_id = get_request_id() + input_rails = GuardrailsAttributes.SPECULATIVE_FIRST_COMPLETED_INPUT_RAILS + generation = GuardrailsAttributes.SPECULATIVE_FIRST_COMPLETED_GENERATION + output_rails = GuardrailsAttributes.SPECULATIVE_FIRST_COMPLETED_OUTPUT_RAILS + none_value = GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE + + released = False + held: list = [] + hold_start: Optional[float] = None + spec_chunks = 0 + + # Human-readable labels for the violation message, keyed by the payload + # ``param``. Keeps tool-result-rail rejections labeled as tool_input_rails + # (matching the non-speculative path) instead of collapsing to input_rails. + reject_labels = {"input_rails": "input rails", "tool_input_rails": "tool input rails"} + + def _mark_reject_input(reason, first_completed, cancellation_event, param): + spec_stats["first_completed"] = first_completed + spec_stats["first_rejector"] = input_rails + spec_stats["cancellation_event"] = cancellation_event + spec_stats["output_rails_wasted_chunks"] = len(held) + spec_stats["safe"] = False + label = reject_labels.get(param, "input rails") + return _frame_for_stream( + self._guardrails_violation_payload(f"Blocked by {label}: {reason}", param), + include_metadata, + ) + + def _mark_release(first_completed): + spec_stats["first_completed"] = first_completed + spec_stats["safe"] = True + spec_stats["release_queue_token_count"] = len(held) + if hold_start is not None: + spec_stats["release_queue_duration_ms"] = (time.monotonic() - hold_start) * 1000 + + try: + async for chunk in base_iterator: + if released: + yield chunk + continue + + # An error chunk during the speculation window is either an + # output-rails violation or a generation error surfaced by the + # base iterator. Either way we short-circuit: cancel the still- + # running input rails task (output-reject early short-circuit) and + # forward the payload. Held (validated-but-unreleased) chunks are + # discarded — the request is being refused. + if _is_stream_error_chunk(chunk): + text = chunk.get("text") if isinstance(chunk, dict) else chunk + is_output_violation = False + if isinstance(text, str): + try: + parsed = json.loads(text) + is_output_violation = ( + isinstance(parsed, dict) and parsed.get("error", {}).get("param") == "output_rails" + ) + except (json.JSONDecodeError, TypeError): + pass + if not input_task.done(): + input_task.cancel() + if is_output_violation: + spec_stats["first_completed"] = output_rails + spec_stats["first_rejector"] = output_rails + spec_stats["output_rails_early_reject"] = True + else: + # Generation error surfaced as a chunk — not a rail rejection. + spec_stats["first_completed"] = generation + spec_stats["first_rejector"] = none_value + spec_stats["cancellation_event"] = GuardrailsAttributes.SPECULATIVE_CANCELLATION_INPUT_RAILS + spec_stats["output_rails_wasted_chunks"] = len(held) + spec_stats["safe"] = False + held.clear() + log.info("[%s] Speculative stream short-circuit (%s)", req_id, spec_stats["first_rejector"]) + yield chunk + return + + spec_chunks += 1 + if hold_start is None: + hold_start = time.monotonic() + held.append(chunk) + + if input_task.done(): + input_result, input_param = input_task.result() + first_completed = input_rails + elif len(held) >= self._speculative_max_buffered_tokens: + # Release buffer full — stop consuming the base iterator and block + # for the verdict (release on pass, teardown on reject). This bounds + # `held`, not the upstream stream queue: the generation task keeps + # producing until the verdict resolves. (To cap total memory instead, + # switch this to abort the request on overflow — cancelling the + # producer stops queue growth at the bound.) + log.info("[%s] Speculative release buffer full (%d); awaiting input verdict", req_id, len(held)) + input_result, input_param = await input_task + first_completed = generation + else: + # Still speculating — keep holding. + continue + + if not input_result.is_safe: + log.info("[%s] Input blocked (speculative streaming): %s", req_id, input_result.reason) + if self._metrics_enabled: + record_request_blocked(RailDirection.INPUT) + refusal = _mark_reject_input( + input_result.reason, + first_completed, + GuardrailsAttributes.SPECULATIVE_CANCELLATION_GENERATION, + input_param, + ) + held.clear() + yield refusal + return + + _mark_release(first_completed) + released = True + for held_chunk in held: + yield held_chunk + held.clear() + + # Stream ended before the input verdict was applied (generation + # finished first, or an empty stream). Await the verdict and either + # flush the held buffer or refuse. + if not released: + input_result, input_param = await input_task + if not input_result.is_safe: + log.info("[%s] Input blocked (speculative streaming, gen-first): %s", req_id, input_result.reason) + if self._metrics_enabled: + record_request_blocked(RailDirection.INPUT) + # Generation already completed, so nothing is cancelled here. + refusal = _mark_reject_input(input_result.reason, generation, none_value, input_param) + held.clear() + yield refusal + return + _mark_release(generation) + for held_chunk in held: + yield held_chunk + held.clear() + finally: + spec_stats["output_rails_speculation_chunks"] = spec_chunks + if not input_task.done(): + input_task.cancel() + with suppress(asyncio.CancelledError, Exception): + await input_task + async def _run_output_rails_in_streaming( self, streaming_handler: AsyncIterator[Union[str, dict]], @@ -1037,6 +1354,7 @@ async def _run_output_rails_in_streaming( *, enabled: Union[bool, list[str]] = True, include_metadata: Optional[bool] = False, + force_check_first: bool = False, ) -> AsyncGenerator[Union[str, dict], None]: """Buffer streamed chunks and run output rails on each batch. @@ -1046,11 +1364,17 @@ async def _run_output_rails_in_streaming( rails. If unsafe, inject an error and stop. - ``stream_first=False``: run output rails first, only yield chunks if safe. + + ``force_check_first`` overrides the configured ``stream_first`` to + check-first behavior. Speculative streaming (SG2) sets this so that + validated chunks are produced (never pre-yielded) — during the + speculation window tokens cannot reach the client before the input + rails verdict is known. """ # Unpack streaming config and get the buffer strategy output_streaming_config = self.config.rails.output.streaming - stream_first = output_streaming_config.stream_first + stream_first = output_streaming_config.stream_first and not force_check_first buffer_strategy = get_buffer_strategy(output_streaming_config) async for chunk_batch in buffer_strategy(streaming_handler): diff --git a/nemoguardrails/guardrails/telemetry.py b/nemoguardrails/guardrails/telemetry.py index 16b22b422d..1d38c61511 100644 --- a/nemoguardrails/guardrails/telemetry.py +++ b/nemoguardrails/guardrails/telemetry.py @@ -382,23 +382,60 @@ def set_speculative_span_attrs( span: Optional["Span"], first_completed: str, first_rejector: str, + *, + rails_duration_ms: Optional[float] = None, + generation_duration_ms: Optional[float] = None, + overlap_ms: Optional[float] = None, + time_saved_ms: Optional[float] = None, + cancellation_event: Optional[str] = None, + output_rails_early_reject: Optional[bool] = None, + output_rails_speculation_chunks: Optional[int] = None, + output_rails_wasted_chunks: Optional[int] = None, + release_queue_duration_ms: Optional[float] = None, + release_queue_token_count: Optional[int] = None, ) -> None: """Stamp speculative-generation outcome attributes on a request span. Records which branch of the speculative race finished first - (input rails vs. main LLM generation) and which one ultimately - rejected the request, on the IORails ``guardrails.request`` span. - Safe to call with ``None`` (no-op) so callers don't have to branch - on whether tracing is enabled — matches the ``record_span_error`` / - ``mark_rail_stop`` idiom. + (input rails, main LLM generation, or output rails) and which one + ultimately rejected the request, on the IORails ``guardrails.request`` + span. Safe to call with ``None`` (no-op) so callers don't have to + branch on whether tracing is enabled — matches the ``record_span_error`` + / ``mark_rail_stop`` idiom. + + The keyword-only timing/streaming signals are optional: the + non-streaming SG1 path leaves them unset, while the streaming SG2 path + supplies the ones it measures. Each is stamped only when provided, so + absent signals never appear as attributes. """ if span is None: return span.set_attribute(GuardrailsAttributes.SPECULATIVE_MODE_ACTIVE, True) span.set_attribute(GuardrailsAttributes.SPECULATIVE_FIRST_COMPLETED, first_completed) span.set_attribute(GuardrailsAttributes.SPECULATIVE_FIRST_REJECTOR, first_rejector) - # TODO: Add it to metrics on next version - # span.set_attribute(GuardrailsAttributes.SPECULATIVE_TIME_SAVED_MS, time_saved_ms) + + if rails_duration_ms is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_RAILS_DURATION_MS, rails_duration_ms) + if generation_duration_ms is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_GENERATION_DURATION_MS, generation_duration_ms) + if overlap_ms is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_OVERLAP_MS, overlap_ms) + if time_saved_ms is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_TIME_SAVED_MS, time_saved_ms) + if cancellation_event is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_CANCELLATION_EVENT, cancellation_event) + if output_rails_early_reject is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_OUTPUT_RAILS_EARLY_REJECT, output_rails_early_reject) + if output_rails_speculation_chunks is not None: + span.set_attribute( + GuardrailsAttributes.SPECULATIVE_OUTPUT_RAILS_SPECULATION_CHUNKS, output_rails_speculation_chunks + ) + if output_rails_wasted_chunks is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_OUTPUT_RAILS_WASTED_CHUNKS, output_rails_wasted_chunks) + if release_queue_duration_ms is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_RELEASE_QUEUE_DURATION_MS, release_queue_duration_ms) + if release_queue_token_count is not None: + span.set_attribute(GuardrailsAttributes.SPECULATIVE_RELEASE_QUEUE_TOKEN_COUNT, release_queue_token_count) # Maps an OpenAI-style ``role`` to the OTEL GenAI legacy event name used diff --git a/nemoguardrails/rails/llm/config.py b/nemoguardrails/rails/llm/config.py index 75efd2e04b..9daecf1403 100644 --- a/nemoguardrails/rails/llm/config.py +++ b/nemoguardrails/rails/llm/config.py @@ -725,8 +725,24 @@ class InputRails(BaseModel): default=False, description=( "If True, input rails run concurrently with LLM generation (speculative execution). " - "Only supported for non-streaming generate_async() calls; stream_async() warns and falls " - "back to sequential execution." + "Supported for both non-streaming generate_async() and streaming stream_async() calls. " + "Streaming speculation only supports check-first behavior: if stream_first is True it is " + "overridden to check-first for speculative requests (with a warning), since tokens cannot " + "reach the client before the input verdict is known." + ), + ) + + speculative_max_buffered_tokens: int = Field( + default=4096, + description=( + "Upper bound on the number of chunks held in the speculative awaiting-release buffer " + "during streaming speculation (chunks approximate tokens). When the bound is reached, " + "the engine pauses speculative output-rail processing and waits for the input rails " + "verdict before buffering more (release on pass, refuse-and-teardown on reject). This " + "bounds only the release buffer — the main LLM keeps generating into the internal stream " + "buffer — so total speculative memory is bounded by the model's finite output and the " + "input-rail latency, not by halting generation. Only used when speculative_generation is " + "True for streaming requests." ), ) diff --git a/nemoguardrails/tracing/constants.py b/nemoguardrails/tracing/constants.py index bc44f66c91..f5e8c17d45 100644 --- a/nemoguardrails/tracing/constants.py +++ b/nemoguardrails/tracing/constants.py @@ -213,6 +213,26 @@ class GuardrailsAttributes: SPECULATIVE_FIRST_COMPLETED_INPUT_RAILS = "input_rails" SPECULATIVE_FIRST_COMPLETED_GENERATION = "generation" + SPECULATIVE_FIRST_COMPLETED_OUTPUT_RAILS = "output_rails" + + # speculative generation timing/outcome attributes + SPECULATIVE_RAILS_DURATION_MS = "speculative_generation.rails_duration_ms" + SPECULATIVE_GENERATION_DURATION_MS = "speculative_generation.generation_duration_ms" + SPECULATIVE_OVERLAP_MS = "speculative_generation.overlap_ms" + SPECULATIVE_TIME_SAVED_MS = "speculative_generation.time_saved_ms" + SPECULATIVE_CANCELLATION_EVENT = "speculative_generation.cancellation_event" + + SPECULATIVE_CANCELLATION_NONE = "none" + SPECULATIVE_CANCELLATION_GENERATION = "generation_cancelled" + SPECULATIVE_CANCELLATION_INPUT_RAILS = "input_rails_cancelled" + SPECULATIVE_CANCELLATION_OUTPUT_RAILS = "output_rails_cancelled" + + # speculative generation streaming-only attributes + SPECULATIVE_OUTPUT_RAILS_EARLY_REJECT = "speculative_generation.output_rails_early_reject" + SPECULATIVE_OUTPUT_RAILS_SPECULATION_CHUNKS = "speculative_generation.output_rails_speculation_chunks" + SPECULATIVE_OUTPUT_RAILS_WASTED_CHUNKS = "speculative_generation.output_rails_wasted_chunks" + SPECULATIVE_RELEASE_QUEUE_DURATION_MS = "speculative_generation.release_queue_duration_ms" + SPECULATIVE_RELEASE_QUEUE_TOKEN_COUNT = "speculative_generation.release_queue_token_count" class SpanNames: diff --git a/tests/guardrails/test_data.py b/tests/guardrails/test_data.py index fbdb25d4bf..2658cfc375 100644 --- a/tests/guardrails/test_data.py +++ b/tests/guardrails/test_data.py @@ -299,3 +299,33 @@ "input": {**NEMOGUARDS_CONFIG["rails"]["input"], "speculative_generation": True}, }, } + +# Speculative generation for streaming: output-rail streaming enabled with +# check-first (stream_first=False), the only mode speculative streaming supports. +NEMOGUARDS_SPECULATIVE_STREAMING_CONFIG = { + **NEMOGUARDS_CONFIG, + "rails": { + **NEMOGUARDS_CONFIG["rails"], + "input": {**NEMOGUARDS_CONFIG["rails"]["input"], "speculative_generation": True}, + "output": { + **NEMOGUARDS_CONFIG["rails"]["output"], + "streaming": { + "enabled": True, + "chunk_size": 3, + "context_size": 1, + "stream_first": False, + }, + }, + }, +} + +# Speculative streaming with input rails only (no output rails) — exercises the +# buffer-and-release path where raw tokens are held until input rails pass. +NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG = { + **NEMOGUARDS_CONFIG, + "rails": { + **NEMOGUARDS_CONFIG["rails"], + "input": {**NEMOGUARDS_CONFIG["rails"]["input"], "speculative_generation": True}, + "output": {"flows": []}, + }, +} diff --git a/tests/guardrails/test_iorails_streaming.py b/tests/guardrails/test_iorails_streaming.py index 211c3c3ac7..a23808acae 100644 --- a/tests/guardrails/test_iorails_streaming.py +++ b/tests/guardrails/test_iorails_streaming.py @@ -77,7 +77,20 @@ def _make_streaming_config(*, enabled: bool = True, stream_first: bool = True) - }, } -_SPECULATIVE_STREAM_WARNING = "speculative_generation is not supported for streaming; falling back to sequential" +_SPECULATIVE_STREAM_FIRST_WARNING = ( + "speculative_generation with stream_first=True is not supported for streaming; " + "forcing check-first behavior for this request" +) + + +def _make_speculative_streaming_config(*, stream_first: bool = False) -> dict: + """NEMOGUARDS_CONFIG with output-rail streaming AND speculative_generation enabled.""" + config = _make_streaming_config(stream_first=stream_first) + config["rails"] = { + **config["rails"], + "input": {**config["rails"]["input"], "speculative_generation": True}, + } + return config async def _mock_stream(model_type, messages, **kwargs): @@ -194,18 +207,32 @@ async def test_include_metadata_allowed_without_output_rails(self, iorails_input assert len(chunks) > 0 @pytest.mark.asyncio - async def test_speculative_generation_streaming_warning_recorded_once(self): - """Default warning filtering records the speculative streaming warning once per call site.""" + async def test_speculative_streaming_runs_no_warning_for_input_only(self): + """Speculative streaming now runs for streaming requests (no fallback warning).""" + async with started_iorails(_INPUT_ONLY_SPECULATIVE_CONFIG) as iorails: + _wire_mocks(iorails) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + chunks = await _collect(iorails.stream_async(messages=[{"role": "user", "content": "hi"}])) + + # Speculation produced streamed content and emitted no stream_first warning + # (input-only config has no output rails, so check-first override is moot). + assert "".join(c for c in chunks if isinstance(c, str)) == "Hello from the streaming LLM! Have a nice day" + assert not [w for w in caught if str(w.message) == _SPECULATIVE_STREAM_FIRST_WARNING] + + @pytest.mark.asyncio + async def test_speculative_stream_first_warns_and_forces_check_first(self): + """speculative_generation + stream_first=True warns once and forces check-first.""" with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): - iorails = IORails(RailsConfig.from_content(config=_INPUT_ONLY_SPECULATIVE_CONFIG)) + iorails = IORails(RailsConfig.from_content(config=_make_speculative_streaming_config(stream_first=True))) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("default") for _ in range(2): iorails.stream_async(messages=[{"role": "user", "content": "hi"}]) - matching_warnings = [warning for warning in caught if str(warning.message) == _SPECULATIVE_STREAM_WARNING] - assert len(matching_warnings) == 1 + matching = [w for w in caught if str(w.message) == _SPECULATIVE_STREAM_FIRST_WARNING] + assert len(matching) == 1 @pytest.mark.asyncio async def test_tools_in_llm_params_forwarded_on_stream_async(self, iorails_input_only): diff --git a/tests/guardrails/test_speculative_generation.py b/tests/guardrails/test_speculative_generation.py index f4e296e771..f829151ad8 100644 --- a/tests/guardrails/test_speculative_generation.py +++ b/tests/guardrails/test_speculative_generation.py @@ -17,6 +17,7 @@ import asyncio import copy +import json import logging from unittest.mock import AsyncMock, patch @@ -28,11 +29,16 @@ from nemoguardrails.guardrails import telemetry from nemoguardrails.guardrails.guardrails_types import RailResult -from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails +from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE, IORails, _is_stream_error_chunk from nemoguardrails.rails.llm.config import RailsConfig -from nemoguardrails.types import LLMResponse +from nemoguardrails.types import LLMResponse, LLMResponseChunk from tests.guardrails.async_helpers import started_iorails -from tests.guardrails.test_data import NEMOGUARDS_CONFIG, NEMOGUARDS_SPECULATIVE_CONFIG +from tests.guardrails.test_data import ( + NEMOGUARDS_CONFIG, + NEMOGUARDS_SPECULATIVE_CONFIG, + NEMOGUARDS_SPECULATIVE_STREAMING_CONFIG, + NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG, +) MESSAGES = [{"role": "user", "content": "hi"}] @@ -471,3 +477,246 @@ async def test_sequential_mode_has_no_speculative_attrs(self, test_tracer, span_ assert "speculative_generation.mode_active" not in attrs assert "speculative_generation.first_completed" not in attrs assert "speculative_generation.first_rejector" not in attrs + + +# ── Streaming speculative generation (SG2) ── + +STREAM_TOKENS = ["Hello", " from", " the", " streaming", " LLM", "!"] + + +def _token_stream(*, delay: float = 0.0, tokens=STREAM_TOKENS, emitted=None): + """Build a mock LLM stream over ``tokens`` with an optional per-token delay. + + ``emitted`` (if given) records each token as it is produced, so tests can + assert whether the generation ran to completion or was aborted mid-stream. + """ + + async def _stream(model_type, messages, **kwargs): + for tok in tokens: + if delay: + await asyncio.sleep(delay) + if emitted is not None: + emitted.append(tok) + yield LLMResponseChunk(delta_content=tok) + + return _stream + + +async def _collect(async_iter): + return [chunk async for chunk in async_iter] + + +def _content_of(chunks): + """Join the plain content chunks, excluding any error/violation payloads.""" + return "".join(c for c in chunks if isinstance(c, str) and not _is_stream_error_chunk(c)) + + +def _error_chunks(chunks): + return [c for c in chunks if isinstance(c, str) and _is_stream_error_chunk(c)] + + +def _wire_stream(iorails, *, input_safe=True, input_delay=0.0, output_safe=True, stream=None): + """Attach input/tool/output rail mocks and a streaming LLM to an IORails instance.""" + + async def _input(messages, *, enabled=True): + if input_delay: + await asyncio.sleep(input_delay) + return RailResult(is_safe=input_safe, reason=None if input_safe else "blocked") + + iorails.rails_manager.are_tool_results_safe = AsyncMock(return_value=RailResult(is_safe=True)) + iorails.rails_manager.is_input_safe = _input + iorails.rails_manager.is_output_safe = AsyncMock( + return_value=RailResult(is_safe=output_safe, reason=None if output_safe else "blocked") + ) + iorails.engine_registry.stream_model_call = stream if stream is not None else _token_stream() + + +@pytest_asyncio.fixture +async def iorails_spec_stream(): + """Speculative streaming with output rails (check-first).""" + async with started_iorails(NEMOGUARDS_SPECULATIVE_STREAMING_CONFIG) as instance: + yield instance + + +@pytest_asyncio.fixture +async def iorails_spec_stream_input_only(): + """Speculative streaming with input rails only (buffer-and-release path).""" + async with started_iorails(NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG) as instance: + yield instance + + +class TestSpeculativeStreaming: + """Streaming speculation (SG2): input rails race the LLM + output rails.""" + + @pytest.mark.asyncio + async def test_streaming_pass_rails_first(self, iorails_spec_stream): + """Fast passing input rails: all validated tokens are delivered, no error.""" + _wire_stream(iorails_spec_stream, stream=_token_stream(delay=0.005)) + chunks = await _collect(iorails_spec_stream.stream_async(MESSAGES)) + assert _content_of(chunks) == "".join(STREAM_TOKENS) + assert not _error_chunks(chunks) + + @pytest.mark.asyncio + async def test_streaming_pass_held_then_released(self, iorails_spec_stream): + """Slow (but passing) input rails: tokens are held during speculation, then flushed.""" + _wire_stream(iorails_spec_stream, input_delay=0.2, stream=_token_stream(delay=0.0)) + chunks = await _collect(iorails_spec_stream.stream_async(MESSAGES)) + assert _content_of(chunks) == "".join(STREAM_TOKENS) + assert not _error_chunks(chunks) + + @pytest.mark.asyncio + async def test_streaming_input_reject_no_leak_and_aborts(self, iorails_spec_stream): + """Input rails reject: no content leaks, refusal emitted, generation aborted mid-stream.""" + emitted = [] + many = [f"tok{i} " for i in range(50)] + _wire_stream( + iorails_spec_stream, + input_safe=False, + input_delay=0.02, + stream=_token_stream(delay=0.01, tokens=many, emitted=emitted), + ) + chunks = await asyncio.wait_for(_collect(iorails_spec_stream.stream_async(MESSAGES)), timeout=3.0) + + assert _content_of(chunks) == "" # nothing leaked before the input verdict + errs = _error_chunks(chunks) + assert errs and json.loads(errs[0])["error"]["param"] == "input_rails" + assert len(emitted) < len(many), "generation should be aborted, not run to completion" + + @pytest.mark.asyncio + async def test_streaming_output_early_reject(self, iorails_spec_stream): + """Output rails reject during speculation: refusal before input verdict, no leak.""" + emitted = [] + many = [f"tok{i} " for i in range(50)] + _wire_stream( + iorails_spec_stream, + input_safe=True, + input_delay=1.0, # input rails still running when output rejects + output_safe=False, + stream=_token_stream(delay=0.01, tokens=many, emitted=emitted), + ) + chunks = await asyncio.wait_for(_collect(iorails_spec_stream.stream_async(MESSAGES)), timeout=3.0) + + assert _content_of(chunks) == "" + errs = _error_chunks(chunks) + assert errs and json.loads(errs[0])["error"]["param"] == "output_rails" + + @pytest.mark.asyncio + async def test_streaming_backpressure_no_drop(self, iorails_spec_stream_input_only): + """A small release buffer applies backpressure without dropping tokens.""" + iorails_spec_stream_input_only._speculative_max_buffered_tokens = 2 + _wire_stream(iorails_spec_stream_input_only, input_delay=0.05, stream=_token_stream(delay=0.0)) + chunks = await _collect(iorails_spec_stream_input_only.stream_async(MESSAGES)) + assert _content_of(chunks) == "".join(STREAM_TOKENS) + + @pytest.mark.asyncio + async def test_streaming_input_only_buffer_and_release(self, iorails_spec_stream_input_only): + """No output rails: raw tokens are held until input passes, then released.""" + _wire_stream(iorails_spec_stream_input_only, input_delay=0.1, stream=_token_stream(delay=0.0)) + chunks = await _collect(iorails_spec_stream_input_only.stream_async(MESSAGES)) + assert _content_of(chunks) == "".join(STREAM_TOKENS) + + @pytest.mark.asyncio + async def test_streaming_input_only_reject_no_leak(self, iorails_spec_stream_input_only): + """No output rails + input reject: no raw tokens leak, refusal emitted.""" + many = [f"t{i} " for i in range(30)] + _wire_stream( + iorails_spec_stream_input_only, + input_safe=False, + input_delay=0.02, + stream=_token_stream(delay=0.01, tokens=many), + ) + chunks = await asyncio.wait_for(_collect(iorails_spec_stream_input_only.stream_async(MESSAGES)), timeout=3.0) + assert _content_of(chunks) == "" + assert _error_chunks(chunks) + + @pytest.mark.asyncio + async def test_streaming_tool_result_reject_uses_tool_input_param(self, iorails_spec_stream_input_only): + """Tool-result-rail rejection surfaces param=tool_input_rails, matching the non-speculative path.""" + io = iorails_spec_stream_input_only + io.rails_manager.are_tool_results_safe = AsyncMock( + return_value=RailResult(is_safe=False, reason="unlinked tool result") + ) + io.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) + io.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + io.engine_registry.stream_model_call = _token_stream(delay=0.01) + + chunks = await asyncio.wait_for(_collect(io.stream_async(MESSAGES)), timeout=3.0) + + assert _content_of(chunks) == "" + errs = _error_chunks(chunks) + assert errs and json.loads(errs[0])["error"]["param"] == "tool_input_rails" + + +def _make_speculative_streaming_tracing_config(): + cfg = copy.deepcopy(NEMOGUARDS_SPECULATIVE_STREAMING_CONFIG) + cfg["tracing"] = {"enabled": True} + return cfg + + +@pytest_asyncio.fixture +async def iorails_spec_stream_tracing(test_tracer): + """Speculative streaming with OTEL tracing, backed by an in-memory exporter.""" + with patch.object(telemetry, "_tracer", test_tracer): + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + config = RailsConfig.from_content(config=_make_speculative_streaming_tracing_config()) + iorails = IORails(config) + async with iorails: + yield iorails + + +class TestSpeculativeStreamingTelemetry: + """Verify speculative span attributes for streaming outcomes.""" + + async def _request_attrs(self, span_exporter): + spans = span_exporter.get_finished_spans() + request_spans = [s for s in spans if s.name == "guardrails.request"] + assert len(request_spans) == 1 + return dict(request_spans[0].attributes) + + @pytest.mark.asyncio + async def test_streaming_pass_span_attrs(self, iorails_spec_stream_tracing, span_exporter): + """Passing stream: mode_active, no rejector, time_saved recorded.""" + _wire_stream(iorails_spec_stream_tracing, stream=_token_stream(delay=0.0)) + await _collect(iorails_spec_stream_tracing.stream_async(MESSAGES)) + + attrs = await self._request_attrs(span_exporter) + assert attrs["speculative_generation.mode_active"] is True + assert attrs["speculative_generation.first_rejector"] == "none" + assert attrs["speculative_generation.output_rails_early_reject"] is False + # both task durations recorded → time_saved present for a safe request + assert attrs["speculative_generation.time_saved_ms"] >= 0.0 + + @pytest.mark.asyncio + async def test_streaming_input_reject_span_attrs(self, iorails_spec_stream_tracing, span_exporter): + """Input reject: first_rejector=input_rails, generation cancelled.""" + many = [f"tok{i} " for i in range(50)] + _wire_stream( + iorails_spec_stream_tracing, + input_safe=False, + input_delay=0.02, + stream=_token_stream(delay=0.01, tokens=many), + ) + await asyncio.wait_for(_collect(iorails_spec_stream_tracing.stream_async(MESSAGES)), timeout=3.0) + + attrs = await self._request_attrs(span_exporter) + assert attrs["speculative_generation.first_rejector"] == "input_rails" + assert attrs["speculative_generation.cancellation_event"] == "generation_cancelled" + + @pytest.mark.asyncio + async def test_streaming_output_early_reject_span_attrs(self, iorails_spec_stream_tracing, span_exporter): + """Output early reject: first_rejector=output_rails, output_rails_early_reject=True.""" + many = [f"tok{i} " for i in range(50)] + _wire_stream( + iorails_spec_stream_tracing, + input_safe=True, + input_delay=1.0, + output_safe=False, + stream=_token_stream(delay=0.01, tokens=many), + ) + await asyncio.wait_for(_collect(iorails_spec_stream_tracing.stream_async(MESSAGES)), timeout=3.0) + + attrs = await self._request_attrs(span_exporter) + assert attrs["speculative_generation.first_rejector"] == "output_rails" + assert attrs["speculative_generation.first_completed"] == "output_rails" + assert attrs["speculative_generation.output_rails_early_reject"] is True + assert attrs["speculative_generation.cancellation_event"] == "input_rails_cancelled" From 0a3b7bc0e951641bc56c9b9fb7d2ee2d63b6d381 Mon Sep 17 00:00:00 2001 From: Oren Hazai Date: Mon, 20 Jul 2026 17:32:29 +0300 Subject: [PATCH 2/5] fix(iorails): log unexpected errors during speculative input-task cleanup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace broad suppress(CancelledError, Exception) in the streaming gate's drain paths with an explicit split — silence expected CancelledError, warn on any other Exception — so a genuine failure in the speculative input-safety task is no longer swallowed silently. --- nemoguardrails/guardrails/iorails.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index c3fc74d61f..ee019e8376 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -1084,8 +1084,16 @@ async def _wrapped_iterator(): if input_task is not None: if not input_task.done(): input_task.cancel() - with suppress(asyncio.CancelledError, Exception): + try: await input_task + except asyncio.CancelledError: + pass + except Exception: + log.warning( + "[%s] Speculative input-safety task raised during defensive cleanup", + req_id, + exc_info=True, + ) except Exception: elapsed_ms = (time.monotonic() - t0) * 1000 log.error("[%s] stream_async failed time=%.1fms", req_id, elapsed_ms, exc_info=True) @@ -1344,8 +1352,16 @@ def _mark_release(first_completed): spec_stats["output_rails_speculation_chunks"] = spec_chunks if not input_task.done(): input_task.cancel() - with suppress(asyncio.CancelledError, Exception): + try: await input_task + except asyncio.CancelledError: + pass + except Exception: + log.warning( + "[%s] Speculative input-safety task raised during cleanup", + req_id, + exc_info=True, + ) async def _run_output_rails_in_streaming( self, From 9c23024cdf32350c40f838bcb97274b0d26dca98 Mon Sep 17 00:00:00 2001 From: Oren Hazai Date: Mon, 20 Jul 2026 17:40:38 +0300 Subject: [PATCH 3/5] fix(iorails): fail closed on unexpected errors in speculative input rails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Catch exceptions from are_tool_results_safe/is_input_safe in _check_speculative_input_safety and convert them to an unsafe RailResult with the correct param, so _gate_on_input surfaces a structured violation payload instead of crashing the stream — matching the non-speculative contract. Add regression tests for both rail families. --- nemoguardrails/guardrails/iorails.py | 20 +++++++++++-- .../guardrails/test_speculative_generation.py | 30 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index ee019e8376..9ad1b2cbb6 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -38,6 +38,7 @@ LLMMessage, LLMMessages, RailDirection, + RailResult, get_request_id, truncate, ) @@ -1167,13 +1168,28 @@ async def _check_speculative_input_safety( path. Runs as its own task so LLM generation can stream concurrently during the speculation window. Records its wall-clock duration into ``spec_stats['rails_duration_ms']`` for telemetry. + + Unexpected exceptions from either safety check fail closed: they are + converted into an unsafe ``RailResult`` so ``_gate_on_input`` surfaces a + structured violation payload instead of crashing the stream, mirroring + the non-speculative ``_generation_task`` error contract. ``CancelledError`` + is not caught, so the gate can still cancel this task on short-circuit. """ + req_id = get_request_id() t0 = time.monotonic() try: - tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) + try: + tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) + except Exception as e: + log.error("[%s] Speculative tool-result rails failed: %s", req_id, e, exc_info=True) + return RailResult(is_safe=False, reason=f"tool input rails error: {e}"), "tool_input_rails" if not tool_result.is_safe: return tool_result, "tool_input_rails" - input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) + try: + input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) + except Exception as e: + log.error("[%s] Speculative input rails failed: %s", req_id, e, exc_info=True) + return RailResult(is_safe=False, reason=f"input rails error: {e}"), "input_rails" return input_result, "input_rails" finally: if spec_stats is not None: diff --git a/tests/guardrails/test_speculative_generation.py b/tests/guardrails/test_speculative_generation.py index f829151ad8..a1b86568f1 100644 --- a/tests/guardrails/test_speculative_generation.py +++ b/tests/guardrails/test_speculative_generation.py @@ -646,6 +646,36 @@ async def test_streaming_tool_result_reject_uses_tool_input_param(self, iorails_ errs = _error_chunks(chunks) assert errs and json.loads(errs[0])["error"]["param"] == "tool_input_rails" + @pytest.mark.asyncio + async def test_streaming_input_rails_exception_fails_closed(self, iorails_spec_stream_input_only): + """An unexpected exception in input rails fails closed as a structured refusal, not a crash.""" + io = iorails_spec_stream_input_only + io.rails_manager.are_tool_results_safe = AsyncMock(return_value=RailResult(is_safe=True)) + io.rails_manager.is_input_safe = AsyncMock(side_effect=RuntimeError("boom")) + io.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + io.engine_registry.stream_model_call = _token_stream(delay=0.01) + + chunks = await asyncio.wait_for(_collect(io.stream_async(MESSAGES)), timeout=3.0) + + assert _content_of(chunks) == "" # nothing leaks on failure + errs = _error_chunks(chunks) + assert errs and json.loads(errs[0])["error"]["param"] == "input_rails" + + @pytest.mark.asyncio + async def test_streaming_tool_rails_exception_fails_closed(self, iorails_spec_stream_input_only): + """An unexpected exception in tool-result rails fails closed with param=tool_input_rails.""" + io = iorails_spec_stream_input_only + io.rails_manager.are_tool_results_safe = AsyncMock(side_effect=RuntimeError("boom")) + io.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) + io.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + io.engine_registry.stream_model_call = _token_stream(delay=0.01) + + chunks = await asyncio.wait_for(_collect(io.stream_async(MESSAGES)), timeout=3.0) + + assert _content_of(chunks) == "" + errs = _error_chunks(chunks) + assert errs and json.loads(errs[0])["error"]["param"] == "tool_input_rails" + def _make_speculative_streaming_tracing_config(): cfg = copy.deepcopy(NEMOGUARDS_SPECULATIVE_STREAMING_CONFIG) From d264a780376a503f15b45d004a18bb453cb324e2 Mon Sep 17 00:00:00 2001 From: Oren Hazai Date: Mon, 20 Jul 2026 18:21:14 +0300 Subject: [PATCH 4/5] Added gt=0 to the speculative_max_buffered_tokens Field, so a non-positive value is now rejected at config-load time instead of silently degrading speculation. --- nemoguardrails/rails/llm/config.py | 1 + tests/test_config_validation.py | 27 +++++++++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/nemoguardrails/rails/llm/config.py b/nemoguardrails/rails/llm/config.py index 9daecf1403..8e257f8922 100644 --- a/nemoguardrails/rails/llm/config.py +++ b/nemoguardrails/rails/llm/config.py @@ -734,6 +734,7 @@ class InputRails(BaseModel): speculative_max_buffered_tokens: int = Field( default=4096, + gt=0, description=( "Upper bound on the number of chunks held in the speculative awaiting-release buffer " "during streaming speculation (chunks approximate tokens). When the bound is reached, " diff --git a/tests/test_config_validation.py b/tests/test_config_validation.py index 7a1ee447df..5d72e5fd26 100644 --- a/tests/test_config_validation.py +++ b/tests/test_config_validation.py @@ -109,6 +109,33 @@ def test_passthrough_and_single_call_incompatibility(): assert "The passthrough mode and the single call dialog" in str(exc_info.value) +def test_speculative_max_buffered_tokens_must_be_positive(): + with pytest.raises(ValueError) as exc_info: + RailsConfig.from_content( + yaml_content=""" + rails: + input: + speculative_generation: True + speculative_max_buffered_tokens: 0 + """, + ) + + assert "speculative_max_buffered_tokens" in str(exc_info.value) + + +def test_speculative_max_buffered_tokens_accepts_positive(): + config = RailsConfig.from_content( + yaml_content=""" + rails: + input: + speculative_generation: True + speculative_max_buffered_tokens: 8 + """, + ) + + assert config.rails.input.speculative_max_buffered_tokens == 8 + + # def test_self_check_facts_prompt_exception(): # with pytest.raises(ValueError) as exc_info: # config = RailsConfig.from_content( From c333503d823a16ac52789168c90004281c11c0f1 Mon Sep 17 00:00:00 2001 From: Oren Hazai Date: Sat, 25 Jul 2026 13:33:40 +0300 Subject: [PATCH 5/5] fix(iorails): correct error classification and telemetry in speculative streaming - Fix AttributeError crash when a chunk parses to JSON with a non-dict error value (e.g. {'error': 'connection refused'}): add _stream_error_field and use it in both _gate_on_input and _run_output_rails_in_streaming, which shared the same defect. - Report rail exceptions as errors rather than policy blocks. A raising rail now yields a generation_error payload and increments requests.errors with span status ERROR, instead of a guardrails_violation that inflated requests.blocked and left the span green. Payload construction is shared with _generation_task via _generation_error_payload. - Rename speculative_max_buffered_tokens to speculative_max_buffered_chunks (it counts chunks) and document that it caps the release buffer only, not total memory. - Omit output_rails_early_reject when no output rails are configured instead of reporting a meaningless False. - Log the stream_first check-first override once at startup via log.warning rather than per request via warnings.warn, which never reaches the logger. - Add tests for error-vs-block classification, no-burst output-rail pacing, both-rails-reject, config-driven release-buffer overflow, and release-queue span attributes. - Update docs: streaming support in the engine feature table, release buffer tuning guidance, the no-output-rails path, and span semantics. --- .../yaml-schema/guardrails-configuration.mdx | 30 ++- docs/observability/tracing/span-reference.mdx | 18 +- docs/reference/engine-feature-support.mdx | 5 +- nemoguardrails/guardrails/iorails.py | 214 +++++++++++++----- nemoguardrails/rails/llm/config.py | 16 +- tests/guardrails/test_data.py | 15 ++ tests/guardrails/test_iorails_streaming.py | 57 +++-- tests/guardrails/test_iorails_telemetry.py | 70 ++++++ .../guardrails/test_speculative_generation.py | 144 ++++++++++-- tests/test_config_validation.py | 12 +- 10 files changed, 460 insertions(+), 121 deletions(-) diff --git a/docs/configure-rails/yaml-schema/guardrails-configuration.mdx b/docs/configure-rails/yaml-schema/guardrails-configuration.mdx index 730c23fa34..560958af51 100644 --- a/docs/configure-rails/yaml-schema/guardrails-configuration.mdx +++ b/docs/configure-rails/yaml-schema/guardrails-configuration.mdx @@ -382,9 +382,35 @@ While the input-rail verdict is still pending (the *speculation window*), valida - **Input rails reject:** discard the held buffer, abort the generation, and return the refusal message. No tokens are leaked. - **Output rails reject during the speculation window:** return the refusal early and cancel the still-running input rails — the request is refused before the input verdict is known. -Streaming speculation only supports check-first behavior. If `rails.output.streaming.stream_first` is `True`, it is overridden to check-first for speculative requests and a warning is emitted. +Streaming speculation only supports check-first behavior. If `rails.output.streaming.stream_first` is `True`, it is overridden to check-first for speculative requests, and a warning is logged at engine startup. -The awaiting-release buffer is bounded by `rails.input.speculative_max_buffered_tokens` (default `4096`, counted in chunks). The buffer fills when the input rails are slower than the main LLM plus output rails, so validated chunks accumulate while the input verdict is still pending. When the bound is reached, the engine pauses speculative processing and waits for the input verdict before buffering more validated chunks (releasing on pass, tearing down on reject). This caps the release buffer only — the main LLM continues generating into the internal stream buffer — so total speculative memory is bounded by the model's finite output and the input-rail latency, not by stopping the backend. +If the main LLM fails during the speculation window, any chunks already held in the buffer are discarded along with the request. Without speculation those chunks would have reached the client before the failure. + +If an input rail *raises* rather than rejecting — for example the safety model is unreachable — the request fails with a `generation_error` payload rather than a `guardrails_violation`, and increments the errors counter rather than the blocked counter. An outage and a refusal are reported as different things. + + + Streaming speculation also runs when no output rails are configured. In that case there is nothing to validate tokens during the speculation window, so raw model tokens are held in the buffer and released once the input rails pass. The client receives no tokens at all until the input verdict, then receives the held tokens at once — expect a pause followed by a burst rather than a smooth stream. Responses on this path are **not** output-rail validated. + + +#### Tuning the Release Buffer + +The awaiting-release buffer is bounded by `rails.input.speculative_max_buffered_chunks` (default `4096`). + +The buffer only grows while the input rails are still pending — that is the speculation window. Once the verdict arrives, the buffer is flushed (on pass) or discarded (on reject) and the stream becomes a pass-through. So its size is simply how much the model produced while the input rails were still working: + +| Input-rail latency | Held at ~50 tokens/sec | +| --- | --- | +| 361 ms | ~20 chunks | +| 5 s | ~250 chunks | +| 80 s | ~4096 chunks (the default bound) | + +At the default the bound is effectively unreachable: a request would time out long before the buffer fills. + +Lowering the value makes the engine stop consuming and wait for the input verdict sooner. With a slow input rail, the client stream pauses at the bound and resumes when the verdict arrives — you trade stream smoothness for a smaller held buffer. + + + This setting does **not** cap total memory. While the gate waits, the main LLM keeps generating into the internal stream buffer, which applies no backpressure to the model. Resident memory is bounded by the maximum response length multiplied by the number of concurrent streams (up to 256), not by this value. Size hosts from expected response length and concurrency rather than from this setting. + ### Configuration Example diff --git a/docs/observability/tracing/span-reference.mdx b/docs/observability/tracing/span-reference.mdx index e2a2a27282..0b3562d837 100644 --- a/docs/observability/tracing/span-reference.mdx +++ b/docs/observability/tracing/span-reference.mdx @@ -223,17 +223,21 @@ When speculative generation is active, the request span also carries these attri | ------------------------------------------------------ | ------- | ------------------------- | --------------------------------------------------------------------------------------------------------------------- | | `speculative_generation.mode_active` | boolean | Speculative generation on | The value `true`. | | `speculative_generation.first_completed` | string | Speculative generation on | The branch that completed first: `input_rails`, `generation`, or `output_rails`. | -| `speculative_generation.first_rejector` | string | Speculative generation on | The branch that rejected the request: `input_rails`, `output_rails`, or `none`. | +| `speculative_generation.first_rejector` | string | Speculative generation on | The branch that rejected the request: `input_rails`, `output_rails`, or `none`. Covers rail *rejections* only — a rail or generation *error* leaves this `none` and marks the span status `ERROR`. | | `speculative_generation.rails_duration_ms` | float | Speculative generation on | Wall-clock time for the input rails. | | `speculative_generation.generation_duration_ms` | float | Speculative generation on | Wall-clock time for the main LLM generation. | | `speculative_generation.overlap_ms` | float | Speculative generation on | Time the input rails and generation ran concurrently. | | `speculative_generation.time_saved_ms` | float | Speculative generation on | Estimated saving versus sequential execution (the overlap for safe requests, `0` for rejected requests). | -| `speculative_generation.cancellation_event` | string | Speculative generation on | Which task was cancelled: `none`, `generation_cancelled`, `input_rails_cancelled`, or `output_rails_cancelled`. | -| `speculative_generation.output_rails_early_reject` | boolean | Streaming speculation | Output rails rejected during the speculation window (short-circuit before the input verdict). | -| `speculative_generation.output_rails_speculation_chunks` | int | Streaming speculation | Number of chunks output rails processed during the speculation window. | -| `speculative_generation.output_rails_wasted_chunks` | int | Streaming speculation | Validated chunks discarded because the request was rejected. | -| `speculative_generation.release_queue_duration_ms` | float | Streaming speculation | Time validated tokens were held in the awaiting-release buffer. | -| `speculative_generation.release_queue_token_count` | int | Streaming speculation | Number of validated tokens held awaiting the input-rail verdict. | +| `speculative_generation.cancellation_event` | string | Speculative generation on | Which task was cancelled: `none`, `generation_cancelled`, `input_rails_cancelled`, or `output_rails_cancelled`. `input_rails_cancelled` covers the output-rails rejection short-circuit as well as error paths. | +| `speculative_generation.output_rails_early_reject` | boolean | Streaming speculation, output rails configured | Output rails rejected during the speculation window (short-circuit before the input verdict). Absent when no output rails are configured. | +| `speculative_generation.output_rails_speculation_chunks` | int | Streaming speculation | Chunks held during the speculation window. With output rails configured these are output-rail-validated; with input rails only they are raw model chunks. | +| `speculative_generation.output_rails_wasted_chunks` | int | Streaming speculation | Held chunks discarded because the request was rejected. | +| `speculative_generation.release_queue_duration_ms` | float | Streaming speculation | Time chunks were held in the awaiting-release buffer. | +| `speculative_generation.release_queue_token_count` | int | Streaming speculation | Number of chunks held awaiting the input-rail verdict. | + + + The attribute prefix is `speculative_generation.`, and the two `output_rails_*` chunk counters describe the release buffer rather than output-rail activity specifically — they are also emitted for input-rails-only configurations. + **`guardrails.rail`** is one span per rail that runs, wrapping the rail's execution. Span kind `INTERNAL`. diff --git a/docs/reference/engine-feature-support.mdx b/docs/reference/engine-feature-support.mdx index ba9795c954..83d6af14b0 100644 --- a/docs/reference/engine-feature-support.mdx +++ b/docs/reference/engine-feature-support.mdx @@ -187,8 +187,9 @@ Both engines run multiple rails in the same direction concurrently when `rails.i For YAML examples, see [Parallel Execution of Input and Output Rails](/configure-guardrails/yaml-schema/guardrails-configuration#parallel-execution-of-input-and-output-rails). `IORails` adds two concurrency capabilities that `LLMRails` does not provide. -Speculative generation (`rails.input.speculative_generation`) runs input rails concurrently with model generation and discards the generation if an input rail blocks, reducing latency on the safe path; it applies to non-streaming generation only. -For a configuration example, see [Speculative Generation](/configure-guardrails/yaml-schema/guardrails-configuration#speculative-generation). +Speculative generation (`rails.input.speculative_generation`) runs input rails concurrently with model generation and discards the generation if an input rail blocks, reducing latency on the safe path; it applies to both non-streaming and streaming requests. +For streaming requests, tokens are held in an in-memory buffer until the input rails pass, and only check-first output-rail behavior is supported. +For a configuration example, see [Speculative Generation](/configure-guardrails/yaml-schema/guardrails-configuration#speculative-generation), and for the streaming specifics see [Streaming Speculative Generation](/configure-guardrails/yaml-schema/guardrails-configuration#streaming-speculative-generation). Admission control through an `AsyncWorkQueue` (and a separate semaphore for streaming) bounds the number of in-flight requests and rejects work when the queue is full. ### Reasoning-Model Support diff --git a/nemoguardrails/guardrails/iorails.py b/nemoguardrails/guardrails/iorails.py index 9ad1b2cbb6..ba4ee48d48 100644 --- a/nemoguardrails/guardrails/iorails.py +++ b/nemoguardrails/guardrails/iorails.py @@ -27,6 +27,7 @@ import warnings from collections.abc import AsyncGenerator, AsyncIterator from contextlib import nullcontext, suppress +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Optional, Union from nemoguardrails.actions.llm.utils import _extract_and_remove_think_tags @@ -109,6 +110,41 @@ def _is_stream_error_chunk(chunk: Union[str, dict]) -> bool: return isinstance(parsed, dict) and "error" in parsed +@dataclass(frozen=True) +class _RailFailure: + """A rail raised, as distinct from a rail rejecting. + + Speculative streaming runs the input rails in their own task, so the gate + receives a verdict rather than an exception. This wrapper keeps an + infrastructure failure distinguishable from an ``is_safe=False`` policy + decision, so the two are reported as ``generation_error`` and + ``guardrails_violation`` respectively rather than collapsing into one. + """ + + exc: BaseException + + +def _stream_error_field(chunk: Union[str, dict], field: str) -> Optional[str]: + """Return ``error.`` from a streamed error payload, or None. + + Returns None for anything that is not an ``{"error": {...}}`` object with + that field — including the case where ``error`` holds a non-dict value, which + ordinary LLM content can produce (e.g. a model emitting + ``{"error": "connection refused"}`` as text). Callers must not assume + ``_is_stream_error_chunk`` guarantees a dict ``error``: it only checks that + the key is present. + """ + text = chunk.get("text") if isinstance(chunk, dict) else chunk + if not isinstance(text, str): + return None + try: + parsed = json.loads(text) + except (json.JSONDecodeError, TypeError): + return None + error = parsed.get("error") if isinstance(parsed, dict) else None + return error.get(field) if isinstance(error, dict) else None + + def _serialize_tool_calls(tool_calls: list[ToolCall]) -> list[dict]: """Serialize ToolCall objects to OpenAI /chat/completions shape. @@ -325,7 +361,22 @@ def __init__(self, config: RailsConfig, *, _report_usage: bool = True) -> None: content_capture_enabled=self._content_capture_enabled, ) self._speculative_generation = config.rails.input.speculative_generation or False - self._speculative_max_buffered_tokens = config.rails.input.speculative_max_buffered_tokens + self._speculative_max_buffered_chunks = config.rails.input.speculative_max_buffered_chunks + + # Streaming speculation supports check-first only: during the speculation + # window tokens cannot reach the client, so a configured stream_first is + # overridden. Decided once here rather than per request so the warning is + # visible at startup instead of only once traffic arrives. + self._speculative_forces_check_first = ( + self._speculative_generation + and self._has_streaming_output_rails + and config.rails.output.streaming.stream_first + ) + if self._speculative_forces_check_first: + log.warning( + "speculative_generation with stream_first=True is not supported for streaming; " + "check-first behavior will be used for speculative streaming requests" + ) # Non-streaming admission queue + worker pool (owned by IORails so # all request-path concurrency controls sit under one roof). The @@ -516,6 +567,18 @@ def _guardrails_violation_payload(message: str, param: str) -> str: } ) + @staticmethod + def _generation_error_payload(message: str) -> str: + """Build the JSON error payload emitted when a streaming request fails. + + Distinct from ``_guardrails_violation_payload``: this is an infrastructure + failure (the request could not be completed), not a policy decision, so + callers must also record error metrics and mark the request span ERROR. + Shared by ``_generation_task`` and the speculative input-rail failure path + so both surface the same ``generation_error`` / ``generation_failed`` shape. + """ + return json.dumps({"error": {"message": message, "type": _GENERATION_ERROR_TYPE, "code": "generation_failed"}}) + async def _do_generate( self, messages: LLMMessages, req_id: str, request_span: Optional["Span"] = None, **kwargs ) -> LLMMessage: @@ -690,7 +753,9 @@ async def _parallel_input_rail_and_response_generation( # Rails passed — wait for generation to finish response = await gen_task - set_speculative_span_attrs(request_span, first_completed, "none") + set_speculative_span_attrs( + request_span, first_completed, GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE + ) else: # Generation finished first — wait for rails verdict response = gen_task.result() @@ -706,7 +771,9 @@ async def _parallel_input_rail_and_response_generation( ) return None - set_speculative_span_attrs(request_span, first_completed, "none") + set_speculative_span_attrs( + request_span, first_completed, GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE + ) log.debug("[%s] Main LLM response: %s", req_id, truncate(response.content)) return response @@ -758,20 +825,12 @@ def stream_async( # Speculative streaming (SG2): input rails race the LLM instead of blocking # before it. Only check-first is supported — during the speculation window # tokens cannot reach the client, so stream_first is overridden to - # check-first for speculative requests. - # - # NOTE: this precondition warning fires per stream_async() call (i.e. at - # request time), not once at engine startup. Flagged here so developers - # know the check-first override is decided per request, not globally. + # check-first for speculative requests. The operator-facing warning for + # that override is emitted once at construction (see __init__), not here: + # warnings.warn() never reaches the logger, so a per-request warning is + # invisible to anyone tailing logs. use_speculative = self._speculative_generation - force_check_first = False - if use_speculative and self._has_streaming_output_rails and self.config.rails.output.streaming.stream_first: - warnings.warn( - "speculative_generation with stream_first=True is not supported for streaming; " - "forcing check-first behavior for this request", - stacklevel=2, - ) - force_check_first = True + force_check_first = self._speculative_forces_check_first if include_metadata and self._has_streaming_output_rails: raise ValueError( @@ -900,9 +959,7 @@ async def _generation_task(request_span, *, run_input_rails: bool = True, spec_s # streaming path. if self._metrics_enabled: record_request_error(e) - error_payload = json.dumps( - {"error": {"message": str(e), "type": _GENERATION_ERROR_TYPE, "code": "generation_failed"}} - ) + error_payload = self._generation_error_payload(str(e)) await streaming_handler.push_chunk(error_payload) await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore[arg-type] finally: @@ -1018,8 +1075,7 @@ async def _wrapped_iterator(): # Gate raw/validated chunks on the input rails verdict during the # speculation window; pass through unchanged when not speculating. - if use_speculative: - assert input_task is not None and spec_stats is not None + if input_task is not None and spec_stats is not None: base_iterator = self._gate_on_input( inner_iterator, input_task, spec_stats, include_metadata=include_metadata ) @@ -1118,6 +1174,16 @@ async def _wrapped_iterator(): # both tasks start together; time_saved is the overlap for # safe requests and 0 for rejected ones. if use_speculative and spec_stats is not None: + # A rail raised inside the speculative input task. + # The gate converted it to an error chunk and + # returned normally, so request_metrics / + # traced_request never saw an exception — record + # the error here, mirroring _generation_task. + rail_error = spec_stats.get("rail_error") + if rail_error is not None: + record_span_error(request_span, rail_error) + if self._metrics_enabled: + record_request_error(rail_error) rails_ms = spec_stats.get("rails_duration_ms") gen_ms = spec_stats.get("generation_duration_ms") overlap_ms = ( @@ -1138,7 +1204,17 @@ async def _wrapped_iterator(): overlap_ms=overlap_ms, time_saved_ms=time_saved_ms, cancellation_event=spec_stats.get("cancellation_event"), - output_rails_early_reject=spec_stats.get("output_rails_early_reject"), + # Only meaningful when output rails are in the + # pipeline: with none configured, stamping False + # would assert they did not early-reject when + # there were none to reject. Gate at the single + # read site — the writes are spread across the + # gate and are easy to miss one of. + output_rails_early_reject=( + spec_stats.get("output_rails_early_reject") + if self._has_streaming_output_rails + else None + ), output_rails_speculation_chunks=spec_stats.get( "output_rails_speculation_chunks" ), @@ -1158,22 +1234,27 @@ async def _check_speculative_input_safety( input_enabled: Union[bool, list[str]] = True, tool_input_enabled: Union[bool, list[str]] = True, spec_stats: Optional[dict] = None, - ): + ) -> tuple[Union[RailResult, _RailFailure], str]: """Concurrent input-safety check for speculative streaming (SG2). - Runs tool-result rails then input rails, returning ``(RailResult, param)`` + Runs tool-result rails then input rails, returning ``(verdict, param)`` for the first failing rail (or the safe input result). ``param`` - identifies the rail family for the client-facing violation payload - (``tool_input_rails`` vs ``input_rails``), matching the non-speculative - path. Runs as its own task so LLM generation can stream concurrently - during the speculation window. Records its wall-clock duration into - ``spec_stats['rails_duration_ms']`` for telemetry. - - Unexpected exceptions from either safety check fail closed: they are - converted into an unsafe ``RailResult`` so ``_gate_on_input`` surfaces a - structured violation payload instead of crashing the stream, mirroring - the non-speculative ``_generation_task`` error contract. ``CancelledError`` - is not caught, so the gate can still cancel this task on short-circuit. + identifies the rail family (``tool_input_rails`` vs ``input_rails``), + matching the non-speculative path. Runs as its own task so LLM + generation can stream concurrently during the speculation window. + Records its wall-clock duration into ``spec_stats['rails_duration_ms']``. + + Unexpected exceptions are wrapped in ``_RailFailure`` rather than + returned as an unsafe ``RailResult``: an infrastructure failure is not a + policy decision, so the gate must be able to tell them apart and emit a + ``generation_error`` payload (plus error metrics and a span ERROR) + instead of a ``content_blocked`` violation. The exception is not + re-raised because ``stream_async`` is an async generator — by this point + the server has already committed a ``200 OK``, so a raise would truncate + the SSE stream instead of delivering a structured error. This matches + ``_generation_task``, the non-speculative streaming path. + ``CancelledError`` is not caught, so the gate can still cancel this task + on short-circuit. """ req_id = get_request_id() t0 = time.monotonic() @@ -1182,14 +1263,14 @@ async def _check_speculative_input_safety( tool_result = await self.rails_manager.are_tool_results_safe(messages, enabled=tool_input_enabled) except Exception as e: log.error("[%s] Speculative tool-result rails failed: %s", req_id, e, exc_info=True) - return RailResult(is_safe=False, reason=f"tool input rails error: {e}"), "tool_input_rails" + return _RailFailure(e), "tool_input_rails" if not tool_result.is_safe: return tool_result, "tool_input_rails" try: input_result = await self.rails_manager.is_input_safe(messages, enabled=input_enabled) except Exception as e: log.error("[%s] Speculative input rails failed: %s", req_id, e, exc_info=True) - return RailResult(is_safe=False, reason=f"input rails error: {e}"), "input_rails" + return _RailFailure(e), "input_rails" return input_result, "input_rails" finally: if spec_stats is not None: @@ -1214,12 +1295,17 @@ async def _gate_on_input( early reject / generation error surfaced as an error chunk) the held buffer is discarded and a refusal / the error payload is emitted. + A rail that *raises* is reported separately from a rail that *rejects*: + the input task returns a ``_RailFailure`` and the gate emits a + ``generation_error`` payload rather than a ``guardrails_violation``, so + callers and metrics can tell an outage from a refusal. + Cancellation (SG2): on input reject the generation task is torn down by the caller's ``finally``; on an output-rails early reject the still- running input rails task is cancelled here so the request aborts before the input verdict arrives. - When the held buffer reaches ``speculative_max_buffered_tokens`` the gate + When the held buffer reaches ``speculative_max_buffered_chunks`` the gate stops consuming the base iterator and blocks on the input verdict, forcing an early resolve (release on pass, teardown on reject). This bounds only the release buffer — the background generation task keeps producing into @@ -1256,6 +1342,22 @@ def _mark_reject_input(reason, first_completed, cancellation_event, param): include_metadata, ) + def _mark_rail_error(exc): + """A rail raised: report an infrastructure failure, not a refusal. + + ``first_rejector`` stays ``none`` because nothing rejected the + request — matching the generation-error branch above. The exception + is stashed for the caller's ``finally``, which owns ``request_span`` + and records the error metric and span status there. + """ + spec_stats["rail_error"] = exc + spec_stats["first_completed"] = input_rails + spec_stats["first_rejector"] = none_value + spec_stats["cancellation_event"] = GuardrailsAttributes.SPECULATIVE_CANCELLATION_GENERATION + spec_stats["output_rails_wasted_chunks"] = len(held) + spec_stats["safe"] = False + return _frame_for_stream(self._generation_error_payload(str(exc)), include_metadata) + def _mark_release(first_completed): spec_stats["first_completed"] = first_completed spec_stats["safe"] = True @@ -1276,16 +1378,7 @@ def _mark_release(first_completed): # forward the payload. Held (validated-but-unreleased) chunks are # discarded — the request is being refused. if _is_stream_error_chunk(chunk): - text = chunk.get("text") if isinstance(chunk, dict) else chunk - is_output_violation = False - if isinstance(text, str): - try: - parsed = json.loads(text) - is_output_violation = ( - isinstance(parsed, dict) and parsed.get("error", {}).get("param") == "output_rails" - ) - except (json.JSONDecodeError, TypeError): - pass + is_output_violation = _stream_error_field(chunk, "param") == "output_rails" if not input_task.done(): input_task.cancel() if is_output_violation: @@ -1312,7 +1405,7 @@ def _mark_release(first_completed): if input_task.done(): input_result, input_param = input_task.result() first_completed = input_rails - elif len(held) >= self._speculative_max_buffered_tokens: + elif len(held) >= self._speculative_max_buffered_chunks: # Release buffer full — stop consuming the base iterator and block # for the verdict (release on pass, teardown on reject). This bounds # `held`, not the upstream stream queue: the generation task keeps @@ -1326,6 +1419,13 @@ def _mark_release(first_completed): # Still speculating — keep holding. continue + if isinstance(input_result, _RailFailure): + log.error("[%s] Input rails failed (speculative streaming)", req_id) + error_chunk = _mark_rail_error(input_result.exc) + held.clear() + yield error_chunk + return + if not input_result.is_safe: log.info("[%s] Input blocked (speculative streaming): %s", req_id, input_result.reason) if self._metrics_enabled: @@ -1351,6 +1451,12 @@ def _mark_release(first_completed): # flush the held buffer or refuse. if not released: input_result, input_param = await input_task + if isinstance(input_result, _RailFailure): + log.error("[%s] Input rails failed (speculative streaming, gen-first)", req_id) + error_chunk = _mark_rail_error(input_result.exc) + held.clear() + yield error_chunk + return if not input_result.is_safe: log.info("[%s] Input blocked (speculative streaming, gen-first): %s", req_id, input_result.reason) if self._metrics_enabled: @@ -1416,13 +1522,9 @@ async def _run_output_rails_in_streaming( # If the batch contains a generation error from _generation_task, # yield it directly and stop — don't feed error JSON through output rails. for chunk in user_output_chunks: - try: - parsed = json.loads(chunk) - if isinstance(parsed, dict) and parsed.get("error", {}).get("type") == _GENERATION_ERROR_TYPE: - yield chunk - return - except (json.JSONDecodeError, TypeError): - pass + if _stream_error_field(chunk, "type") == _GENERATION_ERROR_TYPE: + yield chunk + return if stream_first: for chunk in user_output_chunks: diff --git a/nemoguardrails/rails/llm/config.py b/nemoguardrails/rails/llm/config.py index 8e257f8922..3bde4748ab 100644 --- a/nemoguardrails/rails/llm/config.py +++ b/nemoguardrails/rails/llm/config.py @@ -732,18 +732,18 @@ class InputRails(BaseModel): ), ) - speculative_max_buffered_tokens: int = Field( + speculative_max_buffered_chunks: int = Field( default=4096, gt=0, description=( "Upper bound on the number of chunks held in the speculative awaiting-release buffer " - "during streaming speculation (chunks approximate tokens). When the bound is reached, " - "the engine pauses speculative output-rail processing and waits for the input rails " - "verdict before buffering more (release on pass, refuse-and-teardown on reject). This " - "bounds only the release buffer — the main LLM keeps generating into the internal stream " - "buffer — so total speculative memory is bounded by the model's finite output and the " - "input-rail latency, not by halting generation. Only used when speculative_generation is " - "True for streaming requests." + "during streaming speculation. When the bound is reached, the engine stops consuming " + "and waits for the input rails verdict before buffering more (release on pass, " + "refuse-and-teardown on reject). This caps the release buffer only: the main LLM keeps " + "generating into the internal stream buffer, so it does not cap total memory. The " + "buffer only grows while the input rails are still pending, so at the default it is " + "effectively unreachable — roughly 80 seconds of generation at 50 tokens/sec. Only " + "used when speculative_generation is True for streaming requests." ), ) diff --git a/tests/guardrails/test_data.py b/tests/guardrails/test_data.py index 2658cfc375..268f382023 100644 --- a/tests/guardrails/test_data.py +++ b/tests/guardrails/test_data.py @@ -329,3 +329,18 @@ "output": {"flows": []}, }, } + +# Same, with a release buffer small enough that the overflow path actually runs. +# The 4096 default is ~80s of generation at 50 tok/s, so the backpressure branch is +# unreachable in a test at default settings; setting it through config (rather than +# patching the private attribute) also covers the config -> IORails wiring. +NEMOGUARDS_SPECULATIVE_STREAMING_SMALL_BUFFER_CONFIG = { + **NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG, + "rails": { + **NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG["rails"], + "input": { + **NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG["rails"]["input"], + "speculative_max_buffered_chunks": 2, + }, + }, +} diff --git a/tests/guardrails/test_iorails_streaming.py b/tests/guardrails/test_iorails_streaming.py index a23808acae..e6269df3ac 100644 --- a/tests/guardrails/test_iorails_streaming.py +++ b/tests/guardrails/test_iorails_streaming.py @@ -17,7 +17,7 @@ import asyncio import json -import warnings +import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -77,10 +77,7 @@ def _make_streaming_config(*, enabled: bool = True, stream_first: bool = True) - }, } -_SPECULATIVE_STREAM_FIRST_WARNING = ( - "speculative_generation with stream_first=True is not supported for streaming; " - "forcing check-first behavior for this request" -) +_SPECULATIVE_STREAM_FIRST_LOG = "speculative_generation with stream_first=True is not supported for streaming" def _make_speculative_streaming_config(*, stream_first: bool = False) -> dict: @@ -207,32 +204,44 @@ async def test_include_metadata_allowed_without_output_rails(self, iorails_input assert len(chunks) > 0 @pytest.mark.asyncio - async def test_speculative_streaming_runs_no_warning_for_input_only(self): - """Speculative streaming now runs for streaming requests (no fallback warning).""" + async def test_speculative_streaming_runs_for_input_only(self): + """Speculative streaming runs for streaming requests, with no check-first override. + + An input-only config has no output rails, so ``stream_first`` is moot and + the startup override must not engage. + """ async with started_iorails(_INPUT_ONLY_SPECULATIVE_CONFIG) as iorails: _wire_mocks(iorails) - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - chunks = await _collect(iorails.stream_async(messages=[{"role": "user", "content": "hi"}])) + chunks = await _collect(iorails.stream_async(messages=[{"role": "user", "content": "hi"}])) - # Speculation produced streamed content and emitted no stream_first warning - # (input-only config has no output rails, so check-first override is moot). assert "".join(c for c in chunks if isinstance(c, str)) == "Hello from the streaming LLM! Have a nice day" - assert not [w for w in caught if str(w.message) == _SPECULATIVE_STREAM_FIRST_WARNING] + assert iorails._speculative_forces_check_first is False - @pytest.mark.asyncio - async def test_speculative_stream_first_warns_and_forces_check_first(self): - """speculative_generation + stream_first=True warns once and forces check-first.""" - with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): - iorails = IORails(RailsConfig.from_content(config=_make_speculative_streaming_config(stream_first=True))) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("default") - for _ in range(2): - iorails.stream_async(messages=[{"role": "user", "content": "hi"}]) + def test_speculative_stream_first_logs_at_startup_and_forces_check_first(self, caplog): + """speculative_generation + stream_first=True logs once at construction. - matching = [w for w in caught if str(w.message) == _SPECULATIVE_STREAM_FIRST_WARNING] + Emitted via ``log.warning`` rather than ``warnings.warn``: the latter never + reaches the logger, so an operator tailing logs would never see it. Logging + at construction also surfaces the misconfiguration at startup instead of + only once traffic arrives. + """ + iorails_logger = logging.getLogger("nemoguardrails.guardrails.iorails") + original_propagate = iorails_logger.propagate + iorails_logger.addHandler(caplog.handler) + iorails_logger.propagate = False + try: + with caplog.at_level(logging.WARNING): + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + iorails = IORails( + RailsConfig.from_content(config=_make_speculative_streaming_config(stream_first=True)) + ) + finally: + iorails_logger.removeHandler(caplog.handler) + iorails_logger.propagate = original_propagate + + matching = [r for r in caplog.records if _SPECULATIVE_STREAM_FIRST_LOG in r.message] assert len(matching) == 1 + assert iorails._speculative_forces_check_first is True @pytest.mark.asyncio async def test_tools_in_llm_params_forwarded_on_stream_async(self, iorails_input_only): diff --git a/tests/guardrails/test_iorails_telemetry.py b/tests/guardrails/test_iorails_telemetry.py index 48a68761c6..baf7d54b56 100644 --- a/tests/guardrails/test_iorails_telemetry.py +++ b/tests/guardrails/test_iorails_telemetry.py @@ -2412,3 +2412,73 @@ async def test_json_attrs_on_streaming_llm_span(self, iorails_streaming_content_ span = _main_llm_span(exporter.get_finished_spans()) outputs = json.loads(span.attributes[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES]) assert outputs[0]["parts"][0]["content"] == delivered + + +def _make_speculative_streaming_tracing_config(): + """Input-only streaming + tracing + metrics, with speculation enabled.""" + cfg = copy.deepcopy(_INPUT_ONLY_STREAMING_TRACING_CONFIG) + cfg["rails"]["input"] = {**cfg["rails"]["input"], "speculative_generation": True} + return cfg + + +@pytest_asyncio.fixture +async def iorails_speculative_streaming_tracing(tracer_from_provider): + """Speculative streaming (input rails only) + tracing + metrics enabled.""" + with patch.object(telemetry, "_tracer", tracer_from_provider): + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + config = RailsConfig.from_content(config=_make_speculative_streaming_tracing_config()) + iorails = IORails(config) + async with iorails: + yield iorails + + +class TestSpeculativeStreamingRailErrorTelemetry: + """A rail that raises during speculation is an error, not a block. + + Regression guard for the SG2 review finding: the speculative path + originally converted rail exceptions into ``is_safe=False`` verdicts, which + incremented ``guardrails.requests.blocked``, left ``requests.errors`` at + zero, and left the request span OK — making an outage indistinguishable + from a refusal on every dashboard. + """ + + @pytest.mark.asyncio + async def test_rail_exception_emits_errors_counter_not_blocked( + self, iorails_speculative_streaming_tracing, metric_reader + ): + iorails = iorails_speculative_streaming_tracing + _stub_deep_streaming_pipeline(iorails) + iorails.rails_manager.is_input_safe = AsyncMock(side_effect=RuntimeError("rail down")) + + chunks = [c async for c in iorails.stream_async([{"role": "user", "content": "hi"}])] + assert any(json.loads(c)["error"]["type"] == "generation_error" for c in chunks if c.startswith('{"error"')) + + points = collect_metric_points(metric_reader) + assert points["guardrails.requests.errors"][0].value == 1 + assert points["guardrails.requests.errors"][0].attributes["error.type"] == "RuntimeError" + # The finding itself: a rail outage must never be counted as a block. + assert "guardrails.requests.blocked" not in points + + @pytest.mark.asyncio + async def test_rail_exception_marks_request_span_error(self, iorails_speculative_streaming_tracing, exporter): + iorails = iorails_speculative_streaming_tracing + _stub_deep_streaming_pipeline(iorails) + iorails.rails_manager.is_input_safe = AsyncMock(side_effect=RuntimeError("rail down")) + + [c async for c in iorails.stream_async([{"role": "user", "content": "hi"}])] + + span = _request_span(exporter.get_finished_spans()) + assert span.status.status_code == StatusCode.ERROR + + @pytest.mark.asyncio + async def test_rail_block_still_counts_as_blocked(self, iorails_speculative_streaming_tracing, metric_reader): + """The complement: a genuine refusal keeps the blocked counter.""" + iorails = iorails_speculative_streaming_tracing + _stub_deep_streaming_pipeline(iorails) + iorails.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=False, reason="unsafe")) + + [c async for c in iorails.stream_async([{"role": "user", "content": "bad"}])] + + points = collect_metric_points(metric_reader) + assert points["guardrails.requests.blocked"][0].value == 1 + assert "guardrails.requests.errors" not in points diff --git a/tests/guardrails/test_speculative_generation.py b/tests/guardrails/test_speculative_generation.py index a1b86568f1..995f394fce 100644 --- a/tests/guardrails/test_speculative_generation.py +++ b/tests/guardrails/test_speculative_generation.py @@ -38,6 +38,7 @@ NEMOGUARDS_SPECULATIVE_CONFIG, NEMOGUARDS_SPECULATIVE_STREAMING_CONFIG, NEMOGUARDS_SPECULATIVE_STREAMING_INPUT_ONLY_CONFIG, + NEMOGUARDS_SPECULATIVE_STREAMING_SMALL_BUFFER_CONFIG, ) MESSAGES = [{"role": "user", "content": "hi"}] @@ -545,6 +546,13 @@ async def iorails_spec_stream_input_only(): yield instance +@pytest_asyncio.fixture +async def iorails_spec_stream_small_buffer(): + """Speculative streaming with a 2-chunk release buffer (overflow path).""" + async with started_iorails(NEMOGUARDS_SPECULATIVE_STREAMING_SMALL_BUFFER_CONFIG) as instance: + yield instance + + class TestSpeculativeStreaming: """Streaming speculation (SG2): input rails race the LLM + output rails.""" @@ -601,11 +609,18 @@ async def test_streaming_output_early_reject(self, iorails_spec_stream): assert errs and json.loads(errs[0])["error"]["param"] == "output_rails" @pytest.mark.asyncio - async def test_streaming_backpressure_no_drop(self, iorails_spec_stream_input_only): - """A small release buffer applies backpressure without dropping tokens.""" - iorails_spec_stream_input_only._speculative_max_buffered_tokens = 2 - _wire_stream(iorails_spec_stream_input_only, input_delay=0.05, stream=_token_stream(delay=0.0)) - chunks = await _collect(iorails_spec_stream_input_only.stream_async(MESSAGES)) + async def test_streaming_backpressure_no_drop(self, iorails_spec_stream_small_buffer): + """A small release buffer applies backpressure without dropping tokens. + + The cap comes from config, not a patched private attribute, so this also + covers the ``rails.input.speculative_max_buffered_chunks`` -> ``IORails`` + wiring. + """ + io = iorails_spec_stream_small_buffer + assert io._speculative_max_buffered_chunks == 2, "config value must reach the engine" + + _wire_stream(io, input_delay=0.05, stream=_token_stream(delay=0.0)) + chunks = await _collect(io.stream_async(MESSAGES)) assert _content_of(chunks) == "".join(STREAM_TOKENS) @pytest.mark.asyncio @@ -647,34 +662,111 @@ async def test_streaming_tool_result_reject_uses_tool_input_param(self, iorails_ assert errs and json.loads(errs[0])["error"]["param"] == "tool_input_rails" @pytest.mark.asyncio - async def test_streaming_input_rails_exception_fails_closed(self, iorails_spec_stream_input_only): - """An unexpected exception in input rails fails closed as a structured refusal, not a crash.""" + @pytest.mark.parametrize("failing_rail", ["is_input_safe", "are_tool_results_safe"]) + async def test_streaming_rail_exception_reports_error_not_block(self, iorails_spec_stream_input_only, failing_rail): + """A rail that *raises* is reported as an error, not as a policy refusal. + + The stream must not crash (``stream_async`` is an async generator, so by + this point a server has already committed a 200 OK and a raise would + truncate the SSE response), and it must not masquerade as a block: an + outage is ``generation_error``/``generation_failed``, matching the + non-speculative ``_generation_task`` path, while a refusal stays + ``guardrails_violation``/``content_blocked``. + """ io = iorails_spec_stream_input_only io.rails_manager.are_tool_results_safe = AsyncMock(return_value=RailResult(is_safe=True)) - io.rails_manager.is_input_safe = AsyncMock(side_effect=RuntimeError("boom")) + io.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) io.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) + setattr(io.rails_manager, failing_rail, AsyncMock(side_effect=RuntimeError("boom"))) io.engine_registry.stream_model_call = _token_stream(delay=0.01) chunks = await asyncio.wait_for(_collect(io.stream_async(MESSAGES)), timeout=3.0) assert _content_of(chunks) == "" # nothing leaks on failure errs = _error_chunks(chunks) - assert errs and json.loads(errs[0])["error"]["param"] == "input_rails" + assert errs + error = json.loads(errs[0])["error"] + assert error["type"] == "generation_error" + assert error["code"] == "generation_failed" + assert "boom" in error["message"] + # No ``param``: that field distinguishes rail families on a *block*, and + # this is not one. Matches the _generation_task payload shape exactly. + assert "param" not in error @pytest.mark.asyncio - async def test_streaming_tool_rails_exception_fails_closed(self, iorails_spec_stream_input_only): - """An unexpected exception in tool-result rails fails closed with param=tool_input_rails.""" + async def test_streaming_json_with_string_error_value_does_not_crash(self, iorails_spec_stream_input_only): + """A chunk of ``{"error": "..."}`` must not crash the gate. + + ``_is_stream_error_chunk`` only checks that an ``error`` key exists, so a + non-dict value reaches the output-violation check. Reading ``error.param`` + off a string raised ``AttributeError`` — uncaught, since the handler only + covered JSONDecodeError/TypeError — which escaped ``_gate_on_input`` and + killed the stream. + + The chunk is still forwarded and still ends the stream (that is + ``_is_stream_error_chunk``'s pre-existing behavior for anything shaped like + an error payload); what this locks in is that it is not misclassified as an + *output-rails violation*, and above all that it does not raise. + """ io = iorails_spec_stream_input_only - io.rails_manager.are_tool_results_safe = AsyncMock(side_effect=RuntimeError("boom")) - io.rails_manager.is_input_safe = AsyncMock(return_value=RailResult(is_safe=True)) - io.rails_manager.is_output_safe = AsyncMock(return_value=RailResult(is_safe=True)) - io.engine_registry.stream_model_call = _token_stream(delay=0.01) + content = '{"error": "connection refused"}' + _wire_stream(io, stream=_token_stream(tokens=[content])) chunks = await asyncio.wait_for(_collect(io.stream_async(MESSAGES)), timeout=3.0) + assert [c for c in chunks if isinstance(c, str)] == [content] + + @pytest.mark.asyncio + async def test_streaming_both_rails_reject_first_rejector_wins(self, iorails_spec_stream): + """Unsafe input AND unsafe output: whichever lands first wins, cleanly. + + Output rails run on the first buffered batch while input rails are still + pending, so output rejects first and short-circuits before the input + verdict — one error chunk, no leak, no second refusal appended. + """ + many = [f"tok{i} " for i in range(50)] + _wire_stream( + iorails_spec_stream, + input_safe=False, + input_delay=1.0, + output_safe=False, + stream=_token_stream(delay=0.01, tokens=many), + ) + chunks = await asyncio.wait_for(_collect(iorails_spec_stream.stream_async(MESSAGES)), timeout=3.0) + assert _content_of(chunks) == "" errs = _error_chunks(chunks) - assert errs and json.loads(errs[0])["error"]["param"] == "tool_input_rails" + assert len(errs) == 1, "exactly one rejection should reach the client" + assert json.loads(errs[0])["error"]["param"] == "output_rails" + + @pytest.mark.asyncio + async def test_output_rails_run_at_generation_rate_not_in_a_burst(self, iorails_spec_stream): + """Output-rail calls are spread across generation, not clustered at release. + + This is the property that motivated Option C over buffer-then-release + (design doc S4): output rails sit *upstream* of the hold buffer, so they + process at the LLM's natural rate and the release flushes already-validated + chunks without firing a burst of safety-model calls. Guards against a + refactor that reorders the gate and the output rails. + """ + call_times = [] + + async def _timed_output(messages, chunk, *, enabled=True): + call_times.append(asyncio.get_running_loop().time()) + return RailResult(is_safe=True) + + many = [f"tok{i} " for i in range(30)] + _wire_stream(iorails_spec_stream, input_delay=0.25, stream=_token_stream(delay=0.01, tokens=many)) + iorails_spec_stream.rails_manager.is_output_safe = _timed_output + + await asyncio.wait_for(_collect(iorails_spec_stream.stream_async(MESSAGES)), timeout=5.0) + + assert len(call_times) >= 3, "need several batches to judge the spread" + span = call_times[-1] - call_times[0] + # 30 tokens at 10ms is ~300ms of generation. A burst would collapse every + # call into one tick; spread across even a fraction of that window proves + # the rails ran during generation rather than after the release. + assert span > 0.05, f"output-rail calls clustered in {span:.3f}s — burst regression" def _make_speculative_streaming_tracing_config(): @@ -712,6 +804,7 @@ async def test_streaming_pass_span_attrs(self, iorails_spec_stream_tracing, span attrs = await self._request_attrs(span_exporter) assert attrs["speculative_generation.mode_active"] is True assert attrs["speculative_generation.first_rejector"] == "none" + assert attrs["speculative_generation.first_completed"] == "input_rails" assert attrs["speculative_generation.output_rails_early_reject"] is False # both task durations recorded → time_saved present for a safe request assert attrs["speculative_generation.time_saved_ms"] >= 0.0 @@ -750,3 +843,22 @@ async def test_streaming_output_early_reject_span_attrs(self, iorails_spec_strea assert attrs["speculative_generation.first_completed"] == "output_rails" assert attrs["speculative_generation.output_rails_early_reject"] is True assert attrs["speculative_generation.cancellation_event"] == "input_rails_cancelled" + + @pytest.mark.asyncio + async def test_release_queue_attrs_recorded_when_buffer_fills(self, test_tracer, span_exporter): + """Overflow path: the held buffer's size and hold time are both recorded.""" + cfg = copy.deepcopy(NEMOGUARDS_SPECULATIVE_STREAMING_SMALL_BUFFER_CONFIG) + cfg["tracing"] = {"enabled": True} + with patch.object(telemetry, "_tracer", test_tracer): + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + iorails = IORails(RailsConfig.from_content(config=cfg)) + async with iorails: + _wire_stream(iorails, input_delay=0.05, stream=_token_stream(delay=0.0)) + await asyncio.wait_for(_collect(iorails.stream_async(MESSAGES)), timeout=3.0) + + attrs = await self._request_attrs(span_exporter) + assert attrs["speculative_generation.release_queue_token_count"] == 2 # the configured cap + assert attrs["speculative_generation.release_queue_duration_ms"] > 0.0 + # No output rails configured, so the output-rails-specific flag is absent + # rather than reported as a misleading False. + assert "speculative_generation.output_rails_early_reject" not in attrs diff --git a/tests/test_config_validation.py b/tests/test_config_validation.py index 5d72e5fd26..0a618be373 100644 --- a/tests/test_config_validation.py +++ b/tests/test_config_validation.py @@ -109,31 +109,31 @@ def test_passthrough_and_single_call_incompatibility(): assert "The passthrough mode and the single call dialog" in str(exc_info.value) -def test_speculative_max_buffered_tokens_must_be_positive(): +def test_speculative_max_buffered_chunks_must_be_positive(): with pytest.raises(ValueError) as exc_info: RailsConfig.from_content( yaml_content=""" rails: input: speculative_generation: True - speculative_max_buffered_tokens: 0 + speculative_max_buffered_chunks: 0 """, ) - assert "speculative_max_buffered_tokens" in str(exc_info.value) + assert "speculative_max_buffered_chunks" in str(exc_info.value) -def test_speculative_max_buffered_tokens_accepts_positive(): +def test_speculative_max_buffered_chunks_accepts_positive(): config = RailsConfig.from_content( yaml_content=""" rails: input: speculative_generation: True - speculative_max_buffered_tokens: 8 + speculative_max_buffered_chunks: 8 """, ) - assert config.rails.input.speculative_max_buffered_tokens == 8 + assert config.rails.input.speculative_max_buffered_chunks == 8 # def test_self_check_facts_prompt_exception():