feature/speculative generation m2#2133
Conversation
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
github-actions
Bot
added
status: needs triage
|
Staged Fern docs preview: https://nvidia-preview-pr-2133.docs.buildwithfern.com/nemo/guardrails |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Pouyanpi
added
status: triaged
Greptile SummaryThis PR extends speculative generation to the streaming path (
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/iorails.py | Core of the PR: adds _check_speculative_input_safety, _gate_on_input, and refactors _generation_task to support concurrent input-rail checking during streaming; logic is largely correct but has edge-case issues around exception propagation and exception suppression in cleanup. |
| nemoguardrails/rails/llm/config.py | Adds speculative_max_buffered_tokens: int field; missing lower-bound validator (ge=1) means 0 or negative values produce silent near-serialization behavior rather than an error. |
| nemoguardrails/guardrails/telemetry.py | Extends set_speculative_span_attrs with optional keyword-only timing/streaming parameters; backwards-compatible, clean pattern. |
| nemoguardrails/tracing/constants.py | Adds new GuardrailsAttributes constants for streaming speculation; naming is consistent and well-organized. |
| tests/guardrails/test_speculative_generation.py | Good coverage of streaming pass/reject/output-early-reject/backpressure/input-only/tool-result-param paths plus telemetry span assertions; test timings are generous but still timing-dependent. |
| docs/observability/tracing/span-reference.mdx | Adds span-reference rows for new streaming-speculation attributes; minor inconsistency between 'tokens' in docs and 'chunks' in implementation for release_queue_token_count. |
| tests/guardrails/test_iorails_streaming.py | Replaces the obsolete warn-and-fallback test with two accurate tests for the new behavior; tests are correct. |
| tests/guardrails/test_data.py | Adds two new test config fixtures for speculative streaming scenarios (with and without output rails); straightforward. |
| docs/configure-rails/yaml-schema/guardrails-configuration.mdx | Accurate update describing streaming speculative generation, the check-first override, and backpressure buffer semantics. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Client
participant stream_async
participant input_task as _check_speculative_input_safety
participant gen_task as _generation_task
participant gate as _gate_on_input
participant output_rails as _run_output_rails_in_streaming
participant handler as StreamingHandler
Client->>stream_async: stream_async(messages)
stream_async->>input_task: asyncio.create_task()
stream_async->>gen_task: "asyncio.create_task() run_input_rails=False"
par Concurrent
gen_task->>handler: push_chunk(token) x N
input_task-->>input_task: are_tool_results_safe then is_input_safe
end
loop Speculation window
handler->>output_rails: token batch
output_rails->>gate: validated chunk
gate->>gate: held.append(chunk)
end
alt Input PASS
input_task-->>gate: "is_safe=True"
gate->>Client: flush held buffer then continue streaming
else Input REJECT
input_task-->>gate: "is_safe=False"
gate->>gen_task: cancel
gate->>Client: guardrails_violation no tokens leaked
else Output rails REJECT during speculation
output_rails->>gate: "error chunk param=output_rails"
gate->>input_task: cancel
gate->>Client: guardrails_violation no tokens leaked
else Buffer full backpressure
gate->>input_task: await verdict
input_task-->>gate: result
gate->>Client: flush or reject
end
stream_async->>stream_async: set_speculative_span_attrs
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Client
participant stream_async
participant input_task as _check_speculative_input_safety
participant gen_task as _generation_task
participant gate as _gate_on_input
participant output_rails as _run_output_rails_in_streaming
participant handler as StreamingHandler
Client->>stream_async: stream_async(messages)
stream_async->>input_task: asyncio.create_task()
stream_async->>gen_task: "asyncio.create_task() run_input_rails=False"
par Concurrent
gen_task->>handler: push_chunk(token) x N
input_task-->>input_task: are_tool_results_safe then is_input_safe
end
loop Speculation window
handler->>output_rails: token batch
output_rails->>gate: validated chunk
gate->>gate: held.append(chunk)
end
alt Input PASS
input_task-->>gate: "is_safe=True"
gate->>Client: flush held buffer then continue streaming
else Input REJECT
input_task-->>gate: "is_safe=False"
gate->>gen_task: cancel
gate->>Client: guardrails_violation no tokens leaked
else Output rails REJECT during speculation
output_rails->>gate: "error chunk param=output_rails"
gate->>input_task: cancel
gate->>Client: guardrails_violation no tokens leaked
else Buffer full backpressure
gate->>input_task: await verdict
input_task-->>gate: result
gate->>Client: flush or reject
end
stream_async->>stream_async: set_speculative_span_attrs
Comments Outside Diff (2)
-
nemoguardrails/rails/llm/config.py, line 678-679 (link)The
speculative_max_buffered_tokensfield accepts 0 or negative values. Withmax=0, the conditionlen(held) >= 0fires after appending the very first chunk (sincelen=1 >= 0), causing the gate to immediately block onawait input_taskfor every single token — effectively serializing every request with no speculation benefit. Age=1constraint makes invalid configurations fail at startup rather than silently degrading performance.Prompt To Fix With AI
This is a comment left during a code review. Path: nemoguardrails/rails/llm/config.py Line: 678-679 Comment: The `speculative_max_buffered_tokens` field accepts 0 or negative values. With `max=0`, the condition `len(held) >= 0` fires after appending the very first chunk (since `len=1 >= 0`), causing the gate to immediately block on `await input_task` for every single token — effectively serializing every request with no speculation benefit. A `ge=1` constraint makes invalid configurations fail at startup rather than silently degrading performance. How can I resolve this? If you propose a fix, please make it concise.
-
docs/observability/tracing/span-reference.mdx, line 71 (link)"tokens" vs "chunks" naming inconsistency
The attribute description says "Number of validated tokens held awaiting the input-rail verdict", but
_gate_on_inputsets it tolen(held)whereheldis a list of stream chunks — which may be multi-token batches when output-rail buffering is in play. The config field correctly says "chunks approximate tokens". Consider aligning the docs attribute description to say "chunks" so consumers of this span attribute know what unit they are comparing against.Prompt To Fix With AI
This is a comment left during a code review. Path: docs/observability/tracing/span-reference.mdx Line: 71 Comment: **"tokens" vs "chunks" naming inconsistency** The attribute description says "Number of validated **tokens** held awaiting the input-rail verdict", but `_gate_on_input` sets it to `len(held)` where `held` is a list of stream **chunks** — which may be multi-token batches when output-rail buffering is in play. The config field correctly says "chunks approximate tokens". Consider aligning the docs attribute description to say "chunks" so consumers of this span attribute know what unit they are comparing against. How can I resolve this? If you propose a fix, please make it concise.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Prompt To Fix All With AI
Fix the following 4 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 4
nemoguardrails/rails/llm/config.py:678-679
The `speculative_max_buffered_tokens` field accepts 0 or negative values. With `max=0`, the condition `len(held) >= 0` fires after appending the very first chunk (since `len=1 >= 0`), causing the gate to immediately block on `await input_task` for every single token — effectively serializing every request with no speculation benefit. A `ge=1` constraint makes invalid configurations fail at startup rather than silently degrading performance.
```suggestion
speculative_max_buffered_tokens: int = Field(
default=4096,
ge=1,
```
### Issue 2 of 4
nemoguardrails/guardrails/iorails.py:1288-1290
**Unhandled exception from `input_task.result()` leaves telemetry in an inconsistent state**
If `_check_speculative_input_safety` raises an exception (e.g. a network error in the input-rail backend), `input_task.result()` re-raises it here. The exception propagates out of the `async for` loop, through the `finally` block (which correctly cleans up `input_task`), and then up to the caller as an unhandled stream error. At that point `spec_stats["safe"]` is still `True` and no blocked-request metric has been recorded — so the request silently fails open from a telemetry and metrics standpoint. The same concern applies to the two `await input_task` calls in the backpressure and post-loop paths.
### Issue 3 of 4
nemoguardrails/guardrails/iorails.py:1343-1348
**Overly broad exception suppression in cleanup**
`suppress(asyncio.CancelledError, Exception)` swallows every non-`BaseException` that comes out of `await input_task`. If the task completed with an unexpected exception (e.g. a bug in `_check_speculative_input_safety` that slipped past the `try/finally`), that error would disappear silently here. Consider at minimum logging a warning for unexpected non-`CancelledError` exceptions so they are not completely invisible in production.
### Issue 4 of 4
docs/observability/tracing/span-reference.mdx:71
**"tokens" vs "chunks" naming inconsistency**
The attribute description says "Number of validated **tokens** held awaiting the input-rail verdict", but `_gate_on_input` sets it to `len(held)` where `held` is a list of stream **chunks** — which may be multi-token batches when output-rail buffering is in play. The config field correctly says "chunks approximate tokens". Consider aligning the docs attribute description to say "chunks" so consumers of this span attribute know what unit they are comparing against.
Reviews (1): Last reviewed commit: "Merge branch 'develop' into feature/nemo..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if input_task.done(): | ||
| input_result, input_param = input_task.result() | ||
| first_completed = input_rails |
There was a problem hiding this comment.
input_task.result() leaves telemetry in an inconsistent state
If _check_speculative_input_safety raises an exception (e.g. a network error in the input-rail backend), input_task.result() re-raises it here. The exception propagates out of the async for loop, through the finally block (which correctly cleans up input_task), and then up to the caller as an unhandled stream error. At that point spec_stats["safe"] is still True and no blocked-request metric has been recorded — so the request silently fails open from a telemetry and metrics standpoint. The same concern applies to the two await input_task calls in the backpressure and post-loop paths.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/iorails.py
Line: 1288-1290
Comment:
**Unhandled exception from `input_task.result()` leaves telemetry in an inconsistent state**
If `_check_speculative_input_safety` raises an exception (e.g. a network error in the input-rail backend), `input_task.result()` re-raises it here. The exception propagates out of the `async for` loop, through the `finally` block (which correctly cleans up `input_task`), and then up to the caller as an unhandled stream error. At that point `spec_stats["safe"]` is still `True` and no blocked-request metric has been recorded — so the request silently fails open from a telemetry and metrics standpoint. The same concern applies to the two `await input_task` calls in the backpressure and post-loop paths.
How can I resolve this? If you propose a fix, please make it concise.| 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 |
There was a problem hiding this comment.
suppress(asyncio.CancelledError, Exception) swallows every non-BaseException that comes out of await input_task. If the task completed with an unexpected exception (e.g. a bug in _check_speculative_input_safety that slipped past the try/finally), that error would disappear silently here. Consider at minimum logging a warning for unexpected non-CancelledError exceptions so they are not completely invisible in production.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/iorails.py
Line: 1343-1348
Comment:
**Overly broad exception suppression in cleanup**
`suppress(asyncio.CancelledError, Exception)` swallows every non-`BaseException` that comes out of `await input_task`. If the task completed with an unexpected exception (e.g. a bug in `_check_speculative_input_safety` that slipped past the `try/finally`), that error would disappear silently here. Consider at minimum logging a warning for unexpected non-`CancelledError` exceptions so they are not completely invisible in production.
How can I resolve this? If you propose a fix, please make it concise.
📝 WalkthroughWalkthroughThis PR adds speculative streaming generation (SG2) to IORails, enabling concurrent input-safety checks alongside LLM token streaming with a bounded buffer gate. It expands telemetry attributes and config options, forces check-first behavior for streaming speculation, and updates documentation and tests accordingly. ChangesSpeculative streaming generation (SG2)
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant IORails
participant InputCheckTask
participant GenerationTask
participant Gate as _gate_on_input
Client->>IORails: stream_async(request)
IORails->>InputCheckTask: start _check_speculative_input_safety
IORails->>GenerationTask: start _generation_task(run_input_rails=False)
GenerationTask->>Gate: stream chunks
Gate->>Gate: buffer chunks (up to speculative_max_buffered_tokens)
InputCheckTask-->>Gate: input verdict
alt input passes
Gate->>Client: flush buffered chunks
else input rejected
Gate->>Client: emit refusal/error payload
Gate->>GenerationTask: cancel
end
IORails->>IORails: stamp speculative telemetry on span
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
nemoguardrails/rails/llm/config.py (1)
735-747: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueAdd a lower-bound constraint on
speculative_max_buffered_tokens.A
0or negative value silently degenerates buffering (forces awaiting the verdict on the very first chunk) instead of failing config validation. Considerge=1(or similar) so misconfiguration is caught at load time rather than surfacing as unexpected runtime behavior.♻️ Proposed fix
speculative_max_buffered_tokens: int = Field( default=4096, + ge=1, description=(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemoguardrails/rails/llm/config.py` around lines 735 - 747, The speculative buffer setting in the config model lacks validation for invalid low values, so `speculative_max_buffered_tokens` can be set to 0 or negative and silently break streaming behavior. Add a lower-bound constraint on the `Field` definition in `config.py` for `speculative_max_buffered_tokens` (for example, `ge=1`) so misconfiguration is rejected during config load instead of affecting runtime behavior. Use the existing `Field` declaration for `speculative_max_buffered_tokens` as the place to apply the validation.tests/guardrails/test_speculative_generation.py (1)
604-609: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a timeout guard to the backpressure test.
This test exercises the new "release buffer full →
await input_task" blocking path in_gate_on_input, but unlike the reject-path tests in this same file (which wrap_collect(...)inasyncio.wait_for(..., timeout=3.0)), it has no timeout. A regression in the blocking/backpressure logic (e.g. anawaitthat never resolves) would hang this test in CI instead of failing fast.🧪 Proposed fix
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)) + chunks = await asyncio.wait_for(_collect(iorails_spec_stream_input_only.stream_async(MESSAGES)), timeout=3.0) assert _content_of(chunks) == "".join(STREAM_TOKENS)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/guardrails/test_speculative_generation.py` around lines 604 - 609, The backpressure test in test_streaming_backpressure_no_drop should be wrapped with a timeout guard like the other stream/blocking tests in this file. Update the _collect(iorails_spec_stream_input_only.stream_async(MESSAGES)) call to use asyncio.wait_for with a finite timeout so regressions in _gate_on_input or the release-buffer blocking path fail fast instead of hanging CI.nemoguardrails/guardrails/iorails.py (1)
765-773: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPer-request warning is deduped by Python's default filter, hiding ongoing misconfiguration.
warnings.warn(...)fires on everystream_async()call whenspeculative_generation+stream_first=Trueare misconfigured, but Python's default warnings filter shows identical(message, category, module, lineno)warnings only once per process. After the first request, the misconfiguration silently persists (check-first is still forced correctly, but operators lose visibility into how often it's happening). The testtest_speculative_stream_first_warns_and_forces_check_firstconfirms this: 2 calls yield only 1 caught warning undersimplefilter("default").Consider also emitting via
log.warning(...)(not deduped) so recurring misconfiguration stays visible in production logs, in addition to the one-timewarnings.warn.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemoguardrails/guardrails/iorails.py` around lines 765 - 773, The per-request misconfiguration warning in the streaming path is deduped by Python’s warnings filter, so repeated calls stop surfacing after the first occurrence. Update the `IORails.stream_async` flow where `use_speculative`, `force_check_first`, and `warnings.warn(...)` are handled to also emit a `log.warning(...)` with the same message/context so recurring `speculative_generation` + `stream_first=True` issues remain visible in production, while still keeping the existing warning and forced check-first behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemoguardrails/guardrails/iorails.py`:
- Around line 1343-1348: The cleanup in `iorails.py` around `input_task` is
swallowing all exceptions via `suppress(asyncio.CancelledError, Exception)`,
which hides real failures. Update the `finally` block in `input_task` teardown,
and the matching defensive cleanup in `_wrapped_iterator`, to explicitly await
the task and log any non-`CancelledError` exception instead of discarding it
silently. Follow the pattern already used in `_do_generate_speculative` so
unexpected task errors are preserved for debugging while still allowing
cancellation cleanup.
- Around line 1145-1172: _add exception handling to
`_check_speculative_input_safety` so failures in
`rails_manager.are_tool_results_safe` or `rails_manager.is_input_safe` do not
escape the speculative streaming task. Match the non-speculative contract by
catching unexpected exceptions, converting them into a fail-closed rail/error
result that `_gate_on_input` can surface as a structured
`guardrails_violation`/error payload, and keep the existing
`spec_stats["rails_duration_ms"]` update in the `finally` block. Use the
existing `_check_speculative_input_safety`, `are_tool_results_safe`,
`is_input_safe`, and `_gate_on_input` flow as the integration points.
---
Nitpick comments:
In `@nemoguardrails/guardrails/iorails.py`:
- Around line 765-773: The per-request misconfiguration warning in the streaming
path is deduped by Python’s warnings filter, so repeated calls stop surfacing
after the first occurrence. Update the `IORails.stream_async` flow where
`use_speculative`, `force_check_first`, and `warnings.warn(...)` are handled to
also emit a `log.warning(...)` with the same message/context so recurring
`speculative_generation` + `stream_first=True` issues remain visible in
production, while still keeping the existing warning and forced check-first
behavior.
In `@nemoguardrails/rails/llm/config.py`:
- Around line 735-747: The speculative buffer setting in the config model lacks
validation for invalid low values, so `speculative_max_buffered_tokens` can be
set to 0 or negative and silently break streaming behavior. Add a lower-bound
constraint on the `Field` definition in `config.py` for
`speculative_max_buffered_tokens` (for example, `ge=1`) so misconfiguration is
rejected during config load instead of affecting runtime behavior. Use the
existing `Field` declaration for `speculative_max_buffered_tokens` as the place
to apply the validation.
In `@tests/guardrails/test_speculative_generation.py`:
- Around line 604-609: The backpressure test in
test_streaming_backpressure_no_drop should be wrapped with a timeout guard like
the other stream/blocking tests in this file. Update the
_collect(iorails_spec_stream_input_only.stream_async(MESSAGES)) call to use
asyncio.wait_for with a finite timeout so regressions in _gate_on_input or the
release-buffer blocking path fail fast instead of hanging CI.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 424ff70a-311f-4d77-88f5-5f8947aa5679
📒 Files selected for processing (9)
docs/configure-rails/yaml-schema/guardrails-configuration.mdxdocs/observability/tracing/span-reference.mdxnemoguardrails/guardrails/iorails.pynemoguardrails/guardrails/telemetry.pynemoguardrails/rails/llm/config.pynemoguardrails/tracing/constants.pytests/guardrails/test_data.pytests/guardrails/test_iorails_streaming.pytests/guardrails/test_speculative_generation.py
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Missing exception handling breaks parity with the non-speculative error-handling contract.
_check_speculative_input_safety has no except around are_tool_results_safe / is_input_safe — only a finally for duration. In the non-speculative path, these same calls run inside _generation_task's single broad try/except Exception (lines 819-906), so any internal failure (e.g. a rail model timeout) is caught and converted into a graceful guardrails_violation/error-payload chunk. In the SG2 streaming path these calls were moved into this task, which has no such safety net: an exception here propagates unhandled through _gate_on_input (via input_task.result()/await input_task), breaking the stream ungracefully instead of surfacing a structured error chunk to the client. This is a real regression in exception-handling behavior introduced specifically by toggling speculative_generation on for streaming.
🐛 Proposed fix: fail-closed on unexpected rail exceptions
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,
):
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"
+ except Exception as e:
+ log.error(
+ "[%s] speculative input-safety check failed: %s", get_request_id(), e, exc_info=True
+ )
+ return RailResult(is_safe=False, reason=f"input safety check failed: {e}"), "input_rails"
finally:
if spec_stats is not None:
spec_stats["rails_duration_ms"] = (time.monotonic() - t0) * 1000📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 _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" | |
| except Exception as e: | |
| log.error( | |
| "[%s] speculative input-safety check failed: %s", get_request_id(), e, exc_info=True | |
| ) | |
| return RailResult(is_safe=False, reason=f"input safety check failed: {e}"), "input_rails" | |
| finally: | |
| if spec_stats is not None: | |
| spec_stats["rails_duration_ms"] = (time.monotonic() - t0) * 1000 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemoguardrails/guardrails/iorails.py` around lines 1145 - 1172, _add
exception handling to `_check_speculative_input_safety` so failures in
`rails_manager.are_tool_results_safe` or `rails_manager.is_input_safe` do not
escape the speculative streaming task. Match the non-speculative contract by
catching unexpected exceptions, converting them into a fail-closed rail/error
result that `_gate_on_input` can surface as a structured
`guardrails_violation`/error payload, and keep the existing
`spec_stats["rails_duration_ms"]` update in the `finally` block. Use the
existing `_check_speculative_input_safety`, `are_tool_results_safe`,
`is_input_safe`, and `_gate_on_input` flow as the integration points.
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Cleanup silently swallows all input_task exceptions, hiding real bugs.
with suppress(asyncio.CancelledError, Exception): await input_task discards any exception the task raised, with no logging. The same pattern is duplicated at the defensive teardown in _wrapped_iterator (lines 1084-1088). Compare with _do_generate_speculative's cleanup (lines 638-651), which explicitly logs discarded task errors so they aren't lost during cleanup. If the fix above (adding exception handling to _check_speculative_input_safety) lands, this becomes a pure defense-in-depth gap for other unexpected errors (e.g. cancellation races), but it's still worth logging non-CancelledError exceptions here rather than swallowing them silently.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemoguardrails/guardrails/iorails.py` around lines 1343 - 1348, The cleanup
in `iorails.py` around `input_task` is swallowing all exceptions via
`suppress(asyncio.CancelledError, Exception)`, which hides real failures. Update
the `finally` block in `input_task` teardown, and the matching defensive cleanup
in `_wrapped_iterator`, to explicitly await the task and log any
non-`CancelledError` exception instead of discarding it silently. Follow the
pattern already used in `_do_generate_speculative` so unexpected task errors are
preserved for debugging while still allowing cancellation cleanup.
2 participants
Description
Related Issue(s)
AI Assistance
Summary by CodeRabbit
New Features
Bug Fixes