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: |
tgasser-nv
left a comment
tgasser-nv
left a comment
There was a problem hiding this comment.
Looks good! Just a few high-level comments. Can you rebase on top of develop? Some of the functions in the PR are already merged and available.
Please address the feedback before merging. Let's discuss bounding in-memory held chunks, can be a follow-on PR.
- Can you add tests to get patch line coverage up to 100%? The biggest gap is the "generation finishes before the input check completes" branch is uncovered, there are other smaller gaps.
- Could you add an integration-test description to the PR description? You could use the
nemoguardrails serverornemoguardrails chatapps? This should use a config with live endpoints.
| ) | ||
|
|
||
| @staticmethod | ||
| def _generation_error_payload(message: str) -> str: |
There was a problem hiding this comment.
Can you rebase on top of the latest develop? This helper isn't needed now, there's a build_streaming_error_payload(exc) helper that does the same thing
| 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). |
There was a problem hiding this comment.
nit: This docstring is very verbose, could you trim it down to a handful of lines?
|
|
||
| def _mark_reject_input(reason, first_completed, cancellation_event, param): | ||
| spec_stats["first_completed"] = first_completed | ||
| spec_stats["first_rejector"] = input_rails |
There was a problem hiding this comment.
Does this work for tool-calls which are rejected? Should they have their own category?
| ), | ||
| ) | ||
|
|
||
| speculative_max_buffered_chunks: int = Field( |
There was a problem hiding this comment.
Not blocking for this PR, but let's discuss setting limits on memory-usage. I think we need a second check against the total time which would cancel generation as well as a max chunk count
3 participants
Description
Related Issue(s)
AI Assistance
Summary by CodeRabbit
New Features
Bug Fixes