feat(iorails): add speculative streaming generation - #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 streaming requests (
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/iorails.py | Core streaming speculative generation implementation: adds _check_speculative_input_safety (concurrent input-rail task with exception-to-_RailFailure wrapping), _gate_on_input (bounded hold-and-release buffer with error/reject/release paths), and integrates both into stream_async. All three prior review findings addressed; cancellation, backpressure, and telemetry logic is sound. |
| nemoguardrails/rails/llm/config.py | Adds speculative_max_buffered_chunks field with gt=0 Pydantic constraint. Field name correctly uses chunks terminology. |
| nemoguardrails/tracing/constants.py | Adds new speculative streaming span attribute constants. SPECULATIVE_RELEASE_QUEUE_TOKEN_COUNT uses token in its name but the metric counts stream chunks, inconsistent with speculative_max_buffered_chunks config field. |
| nemoguardrails/guardrails/telemetry.py | Extends set_speculative_span_attrs with all new streaming signals as keyword-only optional params, stamping each only when not None. Clean backward-compatible extension. |
| tests/guardrails/test_speculative_generation.py | Comprehensive streaming speculative generation tests covering pass/reject/output-early-reject/backpressure/input-only/tool-result-param/rail-exception/string-error-value/both-rails-reject/output-rate and telemetry attributes. |
| docs/configure-rails/yaml-schema/guardrails-configuration.mdx | Adds Streaming Speculative Generation section with correct descriptions. One awkward sentence flagged with a suggestion. |
| docs/observability/tracing/span-reference.mdx | Extends speculative generation span-reference table with all new streaming attributes. Descriptions now correctly say chunks throughout. |
| tests/guardrails/test_iorails_streaming.py | Replaces the obsolete warn-and-fallback test with two accurate tests for speculative streaming. |
| tests/guardrails/test_iorails_telemetry.py | Adds TestSpeculativeStreamingRailErrorTelemetry with three regression tests distinguishing rail exceptions from rail rejections in metrics and span status. |
| tests/test_config_validation.py | Adds validation tests for speculative_max_buffered_chunks covering zero rejection and positive acceptance. |
Sequence Diagram
sequenceDiagram
participant Client
participant stream_async
participant input_task as InputTask
participant gen_task as GenerationTask
participant gate as GateOnInput
participant output_rails as OutputRails
stream_async->>input_task: create_task (tool + input rails)
stream_async->>gen_task: "create_task run_input_rails=False"
gen_task->>output_rails: stream LLM chunks
output_rails->>gate: output-rail-validated chunks
gate->>gate: hold chunks in bounded buffer
alt Input rails pass first
input_task-->>gate: "RailResult safe=True"
gate->>Client: flush held buffer then stream normally
else Input rails reject
input_task-->>gate: "RailResult safe=False"
gate->>Client: guardrails_violation payload
stream_async->>gen_task: cancel
else Output rails reject during speculation window
output_rails->>gate: "error chunk param=output_rails"
gate->>input_task: cancel
gate->>Client: guardrails_violation payload early
else Rail raises infrastructure error
input_task-->>gate: _RailFailure wrapping exception
gate->>Client: generation_error payload
stream_async->>gen_task: cancel
else Buffer overflow backpressure
gate->>input_task: await verdict blocking
input_task-->>gate: verdict
gate->>Client: flush on pass or refusal on reject
end
Reviews (7): Last reviewed commit: "Merge branch 'develop' into feature/spec..." | 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.There was a problem hiding this comment.
@greptile-apps Resolved in 9c23024 (not 9cc60c6, edited). That commit wraps both safety calls in _check_speculative_input_safety with try/except Exception, converting an input/tool-rail backend failure into RailResult(is_safe=False) rather than letting it escape the task.
There was a problem hiding this comment.
The commit 9cc60c69c doesn't resolve in this repository, and a code search confirms no except Exception wrapping exists in _check_speculative_input_safety on the current PR head (d836032c). The current code still has only a try/finally (for timing), which means exceptions from are_tool_results_safe or is_input_safe still propagate out of the task unhandled.
Could you double-check that 9cc60c69c was pushed to this branch? It looks like the fix may be on a local commit that hasn't been included in the PR yet.
There was a problem hiding this comment.
@greptile-apps I just pushed and fixed the commit id -- it's 9c23024
| 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.There was a problem hiding this comment.
@greptile-apps -- Resolved in 9c23024. That commit wraps both safety calls in _check_speculative_input_safety with try/except Exception, converting an input/tool-rail backend failure into RailResult(is_safe=False) rather than letting it escape the task.
There was a problem hiding this comment.
That's a clean resolution. Converting backend failures to RailResult(is_safe=False) inside _check_speculative_input_safety means the task itself never completes with an unexpected exception, so the suppress(asyncio.CancelledError, Exception) in the finally block is no longer masking anything — it's only guarding against CancelledError from the explicit input_task.cancel() call. The fail-safe direction (backend error → reject) is also the right default for a guardrails system.
📝 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 | ✅ 6✅ Passed checks (6 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
Pouyanpi
changed the title
…anup 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.
…ails 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.
9cc60c6 to
9c23024
Compare
…itive value is now rejected at config-load time instead of silently degrading speculation.
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
| try: | ||
| parsed = json.loads(text) | ||
| is_output_violation = ( | ||
| isinstance(parsed, dict) and parsed.get("error", {}).get("param") == "output_rails" | ||
| ) | ||
| except (json.JSONDecodeError, TypeError): | ||
| pass |
There was a problem hiding this comment.
error value
parsed.get("error", {}) returns whatever is stored at the "error" key — which may be a plain string (e.g. {"error": "connection refused"}). Calling .get("param") on that string raises AttributeError, which is not caught by except (json.JSONDecodeError, TypeError). The exception escapes _gate_on_input, propagates through the async for chunk in base_iterator loop in _wrapped_iterator, and crashes the stream for what is actually valid LLM content. The fix is to check that the error value is itself a dict before calling .get("param").
| try: | |
| parsed = json.loads(text) | |
| is_output_violation = ( | |
| isinstance(parsed, dict) and parsed.get("error", {}).get("param") == "output_rails" | |
| ) | |
| except (json.JSONDecodeError, TypeError): | |
| pass | |
| try: | |
| parsed = json.loads(text) | |
| error_obj = parsed.get("error") if isinstance(parsed, dict) else None | |
| is_output_violation = ( | |
| isinstance(error_obj, dict) and error_obj.get("param") == "output_rails" | |
| ) | |
| except (json.JSONDecodeError, TypeError): | |
| pass |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/iorails.py
Line: 1282-1288
Comment:
**AttributeError when LLM emits JSON with a non-dict `error` value**
`parsed.get("error", {})` returns whatever is stored at the `"error"` key — which may be a plain string (e.g. `{"error": "connection refused"}`). Calling `.get("param")` on that string raises `AttributeError`, which is not caught by `except (json.JSONDecodeError, TypeError)`. The exception escapes `_gate_on_input`, propagates through the `async for chunk in base_iterator` loop in `_wrapped_iterator`, and crashes the stream for what is actually valid LLM content. The fix is to check that the `error` value is itself a dict before calling `.get("param")`.
```suggestion
try:
parsed = json.loads(text)
error_obj = parsed.get("error") if isinstance(parsed, dict) else None
is_output_violation = (
isinstance(error_obj, dict) and error_obj.get("param") == "output_rails"
)
except (json.JSONDecodeError, TypeError):
pass
```
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
@greptile-apps resolved in latest commit
There was a problem hiding this comment.
The latest commit (c333503) explicitly addresses this — the commit message calls out "Fix AttributeError crash when a chunk parses to JSON with a non-dict error value" and introduces _stream_error_field used in both _gate_on_input and _run_output_rails_in_streaming. Confirmed, thanks!
…ve 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.
Resolve conflicts in the speculative generation path: - iorails.py imports: keep both RailResult (used by the speculative input-safety verdict type) and develop's TimedLLMResponse. - iorails.py rails-pass branch: take develop's timed = await gen_task for the new TimedLLMResponse return, keeping this branch's GuardrailsAttributes.SPECULATIVE_CANCELLATION_NONE constant in place of the 'none' literal. - test_speculative_generation.py imports: keep LLMResponseChunk alongside develop's GenerationResponse and UsageInfo. - test_speculative_generation.py body: both sides appended tests at the same point; keep this branch's streaming speculation suite and develop's TestSpeculativeGenerationTiming. Align _RailFailure with develop's dataclass convention in this area (frozen=True, slots=True, matching RailResult and TimedLLMResponse). Checked that develop's records_out GenerationLog threading does not apply to the speculative streaming path — stream_async captures no records — and that the new RailResult.triggered_rail field is consumed only by check_rails, not by streaming violation payloads.
PR merge guidance@hazai thanks for the PR. GitHub is currently blocking merge for one or more repository requirements:
Relevant guide: |
2 participants
Description
Related Issue(s)
AI Assistance
Summary by CodeRabbit
New Features
Bug Fixes