feat(iorails)!: Support non-streaming GenerationResponse#2178
feat(iorails)!: Support non-streaming GenerationResponse#2178tgasser-nv wants to merge 11 commits into
Conversation
github-actions
Bot
added
size: L
status: needs triage
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
tgasser-nv
changed the title
tgasser-nv
added
status: triaged
…d_actions.llm_calls.{prompt, completion}
tgasser-nv
changed the title
|
@coderabbitai Review this PR |
|
@greptile-apps Review this PR |
|
✅ Action performedReview finished.
|
📝 WalkthroughWalkthroughGeneration APIs now accept typed options. IORails can return structured ChangesStructured generation and rail telemetry
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant IORails
participant RailsManager
participant RailAction
participant EngineRegistry
participant GenerationLog
Client->>IORails: generate(options)
IORails->>RailsManager: run safety rails
RailsManager->>RailAction: execute rail
RailAction->>EngineRegistry: model/API call
EngineRegistry-->>RailAction: response and call metadata
RailsManager-->>IORails: RailResult with RailCallRecord
IORails->>EngineRegistry: main model generation
IORails->>GenerationLog: aggregate records and stats
IORails-->>Client: GenerationResponse
Possibly related PRs
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
@coderabbitai Review this PR |
|
✅ Action performedReview finished.
|
Greptile SummaryThis PR adds
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/iorails.py | Core IORails engine: adds GenerationResponse support for non-streaming generate(); introduces helpers to build GenerationLog, RailCallRecord, and GenerationStats from per-rail records; streaming path silently ignores unsupported options that generate_async() rejects. |
| nemoguardrails/guardrails/guardrails_types.py | Adds RailCallRecord dataclass (frozen, slots) to carry per-rail LLM call data and verdict; adds serialize_prompt helper; no issues found. |
| nemoguardrails/guardrails/rail_action.py | Adds RailLLMCall dataclass and contextvar-based capture in _get_llm_response/_get_api_response; RailAction.run() clears the var at entry to prevent stale data; clean design for parallel-rail safety. |
| nemoguardrails/guardrails/rails_manager.py | Updates _run_rail, _dispatch_tool_call_rail, _run_rails_sequential, and _run_rails_parallel to capture per-rail LLM calls and accumulate RailCallRecords onto RailResult; tool-result multi-exchange flows create only one record (final exchange only). |
| nemoguardrails/guardrails/guardrails.py | Propagates the options parameter through generate() and generate_async() to the underlying engine; clean passthrough with no issues. |
| nemoguardrails/guardrails/engine_registry.py | Adds provider_name() helper to look up the engine/provider string for a model type; used by iorails.py to populate RailCallRecord.llm_provider_name. |
| nemoguardrails/guardrails/actions/content_safety_action.py | Enriches RailResult with structured return_value dict ({allowed, policy_violations}) for GenerationLog; backward-compatible change. |
| nemoguardrails/guardrails/actions/jailbreak_detection_action.py | Adds return_value (the raw is_jailbreak bool) to RailResult for log capture; minimal change, no issues. |
| nemoguardrails/guardrails/actions/topic_safety_action.py | Adds structured return_value ({on_topic}) to RailResult for log capture; clean minimal change. |
| tests/guardrails/test_iorails_generation_response.py | New test file covering GenerationResponse shape, reasoning_content extraction, tool_calls dict shape, and refusal on block; comprehensive unit tests. |
| tests/guardrails/test_iorails_generation_log.py | New test file covering GenerationLog synthesis: activated_rails, llm_calls, stats aggregation, and NotImplementedError for Colang-only fields. |
| tests/guardrails/test_iorails_generation_log_capture.py | New integration-style tests exercising the real RailAction+RailsManager pipeline (mocking only model_call) to verify usage/timing/verdict are threaded through to the GenerationLog. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant C as Caller
participant G as Guardrails.generate_async
participant IO as IORails._do_generate
participant RM as RailsManager
participant RA as RailAction
participant ER as EngineRegistry
participant LB as Log Builder
C->>G: "generate_async(messages, options=GenerationOptions(...))"
G->>IO: _do_generate(messages, options)
IO->>IO: _coerce_generation_options() + _raise_on_unsupported_options()
IO->>RM: is_input_safe(messages)
RM->>RA: action.run(flow, messages)
RA->>RA: _rail_llm_call_var.set(None)
RA->>ER: model_call(model_type, prompt)
ER-->>RA: LLMResponse
RA->>RA: _rail_llm_call_var.set(RailLLMCall(...))
RA-->>RM: RailResult(is_safe, ...)
RM->>RM: get_and_clear_rail_llm_call_contextvar() → RailCallRecord
RM-->>IO: "RailResult(records=(RailCallRecord,))"
IO->>IO: records.extend(input_result.records)
IO->>ER: model_call(main, messages)
ER-->>IO: LLMResponse
IO->>IO: _make_generation_record() → records.append(...)
IO->>RM: is_output_safe(messages, response_text)
RM-->>IO: "RailResult(records=(RailCallRecord,))"
IO->>IO: records.extend(output_result.records)
IO->>LB: _build_generation_log(records, options, duration)
LB->>LB: _activated_rail() + _call_info() + _build_generation_stats()
LB-->>IO: GenerationLog
IO->>IO: _build_generation_response(text, reasoning, response, log)
IO-->>C: GenerationResponse
%%{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 C as Caller
participant G as Guardrails.generate_async
participant IO as IORails._do_generate
participant RM as RailsManager
participant RA as RailAction
participant ER as EngineRegistry
participant LB as Log Builder
C->>G: "generate_async(messages, options=GenerationOptions(...))"
G->>IO: _do_generate(messages, options)
IO->>IO: _coerce_generation_options() + _raise_on_unsupported_options()
IO->>RM: is_input_safe(messages)
RM->>RA: action.run(flow, messages)
RA->>RA: _rail_llm_call_var.set(None)
RA->>ER: model_call(model_type, prompt)
ER-->>RA: LLMResponse
RA->>RA: _rail_llm_call_var.set(RailLLMCall(...))
RA-->>RM: RailResult(is_safe, ...)
RM->>RM: get_and_clear_rail_llm_call_contextvar() → RailCallRecord
RM-->>IO: "RailResult(records=(RailCallRecord,))"
IO->>IO: records.extend(input_result.records)
IO->>ER: model_call(main, messages)
ER-->>IO: LLMResponse
IO->>IO: _make_generation_record() → records.append(...)
IO->>RM: is_output_safe(messages, response_text)
RM-->>IO: "RailResult(records=(RailCallRecord,))"
IO->>IO: records.extend(output_result.records)
IO->>LB: _build_generation_log(records, options, duration)
LB->>LB: _activated_rail() + _call_info() + _build_generation_stats()
LB-->>IO: GenerationLog
IO->>IO: _build_generation_response(text, reasoning, response, log)
IO-->>C: GenerationResponse
Comments Outside Diff (1)
-
nemoguardrails/guardrails/iorails.py, line 1258-1263 (link)Missing unsupported-options validation on streaming path
stream_async()calls_coerce_generation_options()but never calls_raise_on_unsupported_options(), sooptions.state,options.output_vars, andoptions.log.internal_events/colang_historyare silently swallowed instead of raising theValueError/NotImplementedErrorthat identical inputs produce viagenerate_async(). A caller that accidentally passesGenerationOptions(output_vars=["x"])tostream_async()gets no error and no output data, while the same call throughgenerate_async()raises immediately.Prompt To Fix With AI
This is a comment left during a code review. Path: nemoguardrails/guardrails/iorails.py Line: 1258-1263 Comment: **Missing unsupported-options validation on streaming path** `stream_async()` calls `_coerce_generation_options()` but never calls `_raise_on_unsupported_options()`, so `options.state`, `options.output_vars`, and `options.log.internal_events`/`colang_history` are silently swallowed instead of raising the `ValueError`/`NotImplementedError` that identical inputs produce via `generate_async()`. A caller that accidentally passes `GenerationOptions(output_vars=["x"])` to `stream_async()` gets no error and no output data, while the same call through `generate_async()` raises immediately. How can I resolve this? If you propose a fix, please make it concise.
Prompt To Fix All With AI
Fix the following 3 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 3
nemoguardrails/guardrails/iorails.py:1258-1263
**Missing unsupported-options validation on streaming path**
`stream_async()` calls `_coerce_generation_options()` but never calls `_raise_on_unsupported_options()`, so `options.state`, `options.output_vars`, and `options.log.internal_events`/`colang_history` are silently swallowed instead of raising the `ValueError`/`NotImplementedError` that identical inputs produce via `generate_async()`. A caller that accidentally passes `GenerationOptions(output_vars=["x"])` to `stream_async()` gets no error and no output data, while the same call through `generate_async()` raises immediately.
### Issue 2 of 3
nemoguardrails/guardrails/rails_manager.py:385-393
**Tool-result multi-exchange records dropped**
`_run_tool_result_rail` processes exchanges in a loop and then creates a single `RailCallRecord` outside the loop from the final `result`. For a multi-exchange flow that partially succeeds before breaking on an unsafe result, only the breaking exchange is recorded; the safe intermediate exchanges produce no records. If `options.log.activated_rails` is requested, those exchanges are invisible in the `GenerationLog`. This aligns with the current "one record per flow" design, but it is worth confirming this is intentional rather than an oversight now that logging is supported.
### Issue 3 of 3
nemoguardrails/guardrails/iorails.py:280-299
**Speculative path silently drops main-LLM timing from stats**
`_make_generation_record` is called with `started_at=None`, `finished_at=None`, `duration=None` on the speculative path, so `GenerationStats.generation_rails_duration` is always `None` and `llm_calls_duration` excludes the main call when speculative generation is used. This is documented in the code comment, but the PR description and `GenerationResponse` table don't mention it. Users inspecting stats under `options.log` on a speculative-generation config may see unexpectedly empty timing fields.
Reviews (1): Last reviewed commit: "Revert change to nemoguards config main ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| result = await action.run(exchange.results, exchange.calls) | ||
| if not result.is_safe: | ||
| break | ||
| # Tool rails are model-free, so no LLM call is captured (usage stays None). | ||
| call = get_and_clear_rail_llm_call_contextvar() | ||
| result = replace(result, records=(_rail_call_record(flow, "tool_input", result, call),)) | ||
| mark_rail_stop(span, result.is_safe) | ||
| if self._content_capture_enabled: | ||
| all_results = [r for exchange in exchanges for r in exchange.results] |
There was a problem hiding this comment.
_run_tool_result_rail processes exchanges in a loop and then creates a single RailCallRecord outside the loop from the final result. For a multi-exchange flow that partially succeeds before breaking on an unsafe result, only the breaking exchange is recorded; the safe intermediate exchanges produce no records. If options.log.activated_rails is requested, those exchanges are invisible in the GenerationLog. This aligns with the current "one record per flow" design, but it is worth confirming this is intentional rather than an oversight now that logging is supported.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/rails_manager.py
Line: 385-393
Comment:
**Tool-result multi-exchange records dropped**
`_run_tool_result_rail` processes exchanges in a loop and then creates a single `RailCallRecord` outside the loop from the final `result`. For a multi-exchange flow that partially succeeds before breaking on an unsafe result, only the breaking exchange is recorded; the safe intermediate exchanges produce no records. If `options.log.activated_rails` is requested, those exchanges are invisible in the `GenerationLog`. This aligns with the current "one record per flow" design, but it is worth confirming this is intentional rather than an oversight now that logging is supported.
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!
| def _build_generation_stats( | ||
| records: list[RailCallRecord], call_records: list[RailCallRecord], total_duration: Optional[float] | ||
| ) -> GenerationStats: | ||
| """Aggregate per-phase durations and token totals across every recorded call.""" | ||
|
|
||
| def _phase_duration(*rail_types: str) -> Optional[float]: | ||
| total = sum((r.duration or 0.0) for r in records if r.rail_type in rail_types) | ||
| return total or None | ||
|
|
||
| return GenerationStats( | ||
| input_rails_duration=_phase_duration("input", "tool_input"), | ||
| output_rails_duration=_phase_duration("output", "tool_output"), | ||
| generation_rails_duration=_phase_duration("generation"), | ||
| total_duration=total_duration, | ||
| llm_calls_duration=sum((r.duration or 0.0) for r in call_records) or None, | ||
| llm_calls_count=len(call_records), | ||
| llm_calls_total_prompt_tokens=sum(r.usage.input_tokens for r in call_records if r.usage), | ||
| llm_calls_total_completion_tokens=sum(r.usage.output_tokens for r in call_records if r.usage), | ||
| llm_calls_total_tokens=sum(r.usage.total_tokens for r in call_records if r.usage), | ||
| ) |
There was a problem hiding this comment.
_make_generation_record is called with started_at=None, finished_at=None, duration=None on the speculative path, so GenerationStats.generation_rails_duration is always None and llm_calls_duration excludes the main call when speculative generation is used. This is documented in the code comment, but the PR description and GenerationResponse table don't mention it. Users inspecting stats under options.log on a speculative-generation config may see unexpectedly empty timing fields.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/iorails.py
Line: 280-299
Comment:
**Speculative path silently drops main-LLM timing from stats**
`_make_generation_record` is called with `started_at=None`, `finished_at=None`, `duration=None` on the speculative path, so `GenerationStats.generation_rails_duration` is always `None` and `llm_calls_duration` excludes the main call when speculative generation is used. This is documented in the code comment, but the PR description and `GenerationResponse` table don't mention it. Users inspecting stats under `options.log` on a speculative-generation config may see unexpectedly empty timing fields.
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!
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
nemoguardrails/guardrails/iorails.py (1)
958-1067: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftCapture speculative generation before applying the input-rail verdict.
When generation finishes first and input rails later block, Lines 1051-1058 return without appending the completed main call at Line 1067. Simultaneous completion in the rails-first branch has the same omission. This undercounts calls, tokens, and cost on blocked requests; successful speculative calls also lose timing. Time and record the generation task at its execution boundary, regardless of whether its response is ultimately returned.
🤖 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 958 - 1067, Update _parallel_input_rail_and_response_generation to capture and record the completed generation task before applying the input-rail verdict, including both generation-first and simultaneous-completion paths. Time the main LLM call at its execution boundary and append its generation record with usage, model, timestamps, duration, provider, and main_prompt even when input rails later reject the response; avoid recording the call twice.nemoguardrails/guardrails/rails_manager.py (1)
438-461: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winProcess every completed task before returning an unsafe result.
asyncio.waitcan place multiple tasks indone. If the first sorted result is unsafe, Line 460 returns before collecting records—or retrieving exceptions—from the remaining completed tasks. This violates the “every rail that ran” log contract. Consume the fulldonebatch, select the first unsafe result, then cancel only pending tasks.🤖 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/rails_manager.py` around lines 438 - 461, Update the task-processing loop in the rail execution method to process every task in each `done` batch, collecting records and retrieving results or exceptions before deciding to return. Track the first unsafe result encountered in sorted task order, then cancel and await only `pending_tasks`; return that unsafe result with all collected records after the batch is fully consumed.
🧹 Nitpick comments (1)
tests/guardrails/test_guardrails.py (1)
126-127: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a non-
Noneoptions forwarding test.These assertions only cover the default path. Add one sync/async wrapper test with a sentinel
GenerationOptionsor dict and assert that the same object is passed to the underlying engine, preventing regressions that accept but drop the new argument.🤖 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_guardrails.py` around lines 126 - 127, Add a test covering non-None options for both sync and async guardrail wrappers, using a sentinel GenerationOptions instance or dict. Invoke the wrappers with that same object and assert rails_engine.generate and generate_async each receive it unchanged alongside messages, while preserving the existing default-path assertions.
🤖 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/guardrails_types.py`:
- Around line 134-141: Update serialize_prompt so captured prompts preserve all
message metadata, including tool_calls, tool_call_id, names, and reasoning-only
fields, instead of formatting only role and content. Reuse the existing
canonical message serialization if available; otherwise encode each complete
message structure while retaining the established prompt ordering and
separation. Add coverage for tool-call and reasoning-only turns.
In `@nemoguardrails/guardrails/guardrails.py`:
- Around line 221-255: The generate_async overloads currently share identical
parameters, causing type checkers to select only the first str return type.
Update generate_async so overloads distinguish calls by the options value, or
replace them with a single accurate union return annotation covering str, dict,
GenerationResponse, and tuple[dict, dict].
In `@nemoguardrails/guardrails/rail_action.py`:
- Around line 197-221: Update both LLM-call helpers surrounding
EngineRegistry.model_call and the corresponding helper at the referenced later
section to record attempted calls when the await raises. In each exception path,
capture elapsed timing and populate _rail_llm_call_var with the available
request/model/provider and failure-safe metadata, then re-raise the original
exception so existing error handling remains unchanged. Preserve the current
successful-response recording behavior.
---
Outside diff comments:
In `@nemoguardrails/guardrails/iorails.py`:
- Around line 958-1067: Update _parallel_input_rail_and_response_generation to
capture and record the completed generation task before applying the input-rail
verdict, including both generation-first and simultaneous-completion paths. Time
the main LLM call at its execution boundary and append its generation record
with usage, model, timestamps, duration, provider, and main_prompt even when
input rails later reject the response; avoid recording the call twice.
In `@nemoguardrails/guardrails/rails_manager.py`:
- Around line 438-461: Update the task-processing loop in the rail execution
method to process every task in each `done` batch, collecting records and
retrieving results or exceptions before deciding to return. Track the first
unsafe result encountered in sorted task order, then cancel and await only
`pending_tasks`; return that unsafe result with all collected records after the
batch is fully consumed.
---
Nitpick comments:
In `@tests/guardrails/test_guardrails.py`:
- Around line 126-127: Add a test covering non-None options for both sync and
async guardrail wrappers, using a sentinel GenerationOptions instance or dict.
Invoke the wrappers with that same object and assert rails_engine.generate and
generate_async each receive it unchanged alongside messages, while preserving
the existing default-path assertions.
🪄 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: 0f428c31-8892-4e50-a8d1-5398aa3ee1be
📒 Files selected for processing (19)
nemoguardrails/base_guardrails.pynemoguardrails/guardrails/actions/content_safety_action.pynemoguardrails/guardrails/actions/jailbreak_detection_action.pynemoguardrails/guardrails/actions/topic_safety_action.pynemoguardrails/guardrails/engine_registry.pynemoguardrails/guardrails/guardrails.pynemoguardrails/guardrails/guardrails_types.pynemoguardrails/guardrails/iorails.pynemoguardrails/guardrails/rail_action.pynemoguardrails/guardrails/rails_manager.pytests/guardrails/test_guardrails.pytests/guardrails/test_iorails.pytests/guardrails/test_iorails_generation_log.pytests/guardrails/test_iorails_generation_log_capture.pytests/guardrails/test_iorails_generation_response.pytests/guardrails/test_jailbreak_detection_iorails_actions.pytests/guardrails/test_rails_manager.pytests/guardrails/test_tool_rails_iorails.pytests/guardrails/test_topic_safety_iorails_actions.py
| def serialize_prompt(messages: list[dict]) -> str: | ||
| """Render a chat message list to a role-labeled string for GenerationLog's ``prompt``. | ||
| Content parity with LLMRails' logged prompt, not byte-for-byte format parity: each | ||
| message becomes ``"<role>: <content>"`` and messages are blank-line separated. A | ||
| message with no content (e.g. a reasoning-only or tool-call turn) renders as empty. | ||
| """ | ||
| return "\n\n".join(f"{m.get('role', '')}: {m.get('content') or ''}" for m in messages) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Preserve non-text fields in captured prompts.
Line 141 discards fields such as tool_calls, tool_call_id, names, and reasoning-only metadata, so GenerationLog.llm_calls[*].prompt can misrepresent the actual model request. Use the existing canonical message serialization or encode all message fields, with coverage for tool-call and reasoning-only turns.
As per coding guidelines, preserve non-text model metadata and do not drop reasoning-only content.
🤖 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/guardrails_types.py` around lines 134 - 141, Update
serialize_prompt so captured prompts preserve all message metadata, including
tool_calls, tool_call_id, names, and reasoning-only fields, instead of
formatting only role and content. Reuse the existing canonical message
serialization if available; otherwise encode each complete message structure
while retaining the established prompt ordering and separation. Add coverage for
tool-call and reasoning-only turns.
Source: Coding guidelines
| @overload | ||
| async def generate_async(self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs) -> str: ... | ||
| async def generate_async( | ||
| self, | ||
| prompt: str | None = None, | ||
| messages: LLMMessages | None = None, | ||
| options: Optional[Union[dict, GenerationOptions]] = None, | ||
| **kwargs, | ||
| ) -> str: ... | ||
|
|
||
| @overload | ||
| async def generate_async( | ||
| self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs | ||
| self, | ||
| prompt: str | None = None, | ||
| messages: LLMMessages | None = None, | ||
| options: Optional[Union[dict, GenerationOptions]] = None, | ||
| **kwargs, | ||
| ) -> dict: ... | ||
|
|
||
| @overload | ||
| async def generate_async( | ||
| self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs | ||
| self, | ||
| prompt: str | None = None, | ||
| messages: LLMMessages | None = None, | ||
| options: Optional[Union[dict, GenerationOptions]] = None, | ||
| **kwargs, | ||
| ) -> GenerationResponse: ... | ||
|
|
||
| @overload | ||
| async def generate_async( | ||
| self, prompt: str | None = None, messages: LLMMessages | None = None, **kwargs | ||
| self, | ||
| prompt: str | None = None, | ||
| messages: LLMMessages | None = None, | ||
| options: Optional[Union[dict, GenerationOptions]] = None, | ||
| **kwargs, | ||
| ) -> tuple[dict, dict]: ... |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant method and nearby definitions.
file="nemoguardrails/guardrails/guardrails.py"
wc -l "$file"
sed -n '180,310p' "$file"
# Find the concrete implementation and any related overloads in the repo.
rg -n "`@overload`|def generate_async|def generate\(" nemoguardrails/guardrails/guardrails.py nemoguardrails -g '*.py'Repository: NVIDIA-NeMo/Guardrails
Length of output: 7333
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the underlying LLMRails overloads and any public type declarations.
sed -n '1340,1415p' nemoguardrails/rails/llm/llmrails.py
sed -n '1415,1515p' nemoguardrails/rails/llm/llmrails.py
sed -n '220,290p' nemoguardrails/types.py
sed -n '1,90p' nemoguardrails/base_guardrails.pyRepository: NVIDIA-NeMo/Guardrails
Length of output: 11942
Make the async overloads distinguish options or collapse them into one union return type.
All four generate_async overloads take the same parameters, so type checkers will always pick the first str overload and lose the GenerationResponse/dict/tuple return types. Split the options=None case from the structured-response case, or replace the overload set with one accurate union return annotation.
🤖 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/guardrails.py` around lines 221 - 255, The
generate_async overloads currently share identical parameters, causing type
checkers to select only the first str return type. Update generate_async so
overloads distinguish calls by the options value, or replace them with a single
accurate union return annotation covering str, dict, GenerationResponse, and
tuple[dict, dict].
| """Call an LLM via EngineRegistry and return the structured response. | ||
|
|
||
| Captures usage/model/timing into a request-scoped contextvar so RailsManager can | ||
| record this call in the GenerationLog after the rail runs. | ||
| """ | ||
| if not model_type: | ||
| raise RuntimeError("model_type is required for LLM calls") | ||
| return await self.engine_registry.model_call(model_type, messages, **kwargs) | ||
| started_at = time.time() | ||
| t0 = time.monotonic() | ||
| response = await self.engine_registry.model_call(model_type, messages, **kwargs) | ||
| duration = time.monotonic() - t0 | ||
| finished_at = time.time() | ||
| _rail_llm_call_var.set( | ||
| RailLLMCall( | ||
| usage=response.usage, | ||
| llm_model_name=response.model, | ||
| request_id=response.request_id, | ||
| provider_name=self.engine_registry.provider_name(model_type), | ||
| prompt=serialize_prompt(messages), | ||
| completion=response.content, | ||
| started_at=started_at, | ||
| finished_at=finished_at, | ||
| duration=duration, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Record attempted model and API calls when they fail.
Both helpers populate _rail_llm_call_var only after the await succeeds. If the external call raises, run() returns an unsafe result but RailsManager records made_call=False, causing the generated log and call count to claim no call occurred. Capture the attempt and elapsed timing in an exception path before re-raising.
Also applies to: 224-254
🤖 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/rail_action.py` around lines 197 - 221, Update both
LLM-call helpers surrounding EngineRegistry.model_call and the corresponding
helper at the referenced later section to record attempted calls when the await
raises. In each exception path, capture elapsed timing and populate
_rail_llm_call_var with the available request/model/provider and failure-safe
metadata, then re-raise the original exception so existing error handling
remains unchanged. Preserve the current successful-response recording behavior.
Greptile SummaryThis PR wires up the
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/iorails.py | Core change: adds GenerationResponse support to generate_async/generate; introduces record collection pipeline, per-phase log synthesis, and refusal shaping for both return contracts. The speculative path correctly omits per-call timing (documented) but this causes llm_calls_duration to be underreported. |
| nemoguardrails/guardrails/rail_action.py | Adds RailLLMCall dataclass and _rail_llm_call_var ContextVar to capture LLM/API call metadata per rail; ContextVar is cleared at the start of each rail run and set after the call, correctly isolated per asyncio Task for parallel rails. |
| nemoguardrails/guardrails/rails_manager.py | Adds _rail_call_record builder and wires record collection into _run_rail, _run_tool_call_rail, _run_tool_result_rail, and both sequential/parallel runners; records propagate correctly through collected lists and are merged into the final RailResult. |
| nemoguardrails/guardrails/guardrails_types.py | Introduces frozen RailCallRecord dataclass; records/return_value on RailResult are correctly excluded from equality/hashing via compare=False; serialize_prompt helper is straightforward. |
| nemoguardrails/guardrails/guardrails.py | Promotes options from **kwargs passthrough to an explicit named parameter on generate/generate_async; overload signatures and delegation to rails_engine are updated consistently. |
| nemoguardrails/guardrails/actions/jailbreak_detection_action.py | Adds return_value to RailResult, but uses a raw bool (the jailbreak field value) rather than a structured dict like content_safety ({allowed, policy_violations}) and topic_safety ({on_topic}), creating a type inconsistency in ExecutedAction.return_value across rail types. |
| nemoguardrails/guardrails/actions/content_safety_action.py | Consistently adds structured return_value dicts ({allowed, policy_violations}) for both safe and unsafe outcomes. |
| nemoguardrails/guardrails/actions/topic_safety_action.py | Adds structured {on_topic: bool} return_value to both safe and unsafe RailResult paths. |
| nemoguardrails/guardrails/engine_registry.py | Adds provider_name() helper that delegates to _get_engine; straightforward and correctly typed. |
| tests/guardrails/test_iorails_generation_response.py | New test file; good coverage of the structured vs bare return trigger, response/reasoning/tool_call fields, refusal shaping, and unsupported-option validation (state, output_vars, internal_events). |
| tests/guardrails/test_iorails_generation_log.py | New test file covering GenerationLog synthesis: activated_rails, llm_calls, stats, token aggregation, and blocking-rail partial logs. |
| tests/guardrails/test_iorails_generation_log_capture.py | New test file exercising log capture with real RailCallRecord construction through the full pipeline. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant C as Caller
participant GR as Guardrails.generate_async
participant IO as IORails._do_generate
participant RM as RailsManager
participant RA as RailAction._get_llm_response
participant CV as _rail_llm_call_var (ContextVar)
participant ER as EngineRegistry
C->>GR: generate_async(messages, options)
GR->>IO: "_do_generate(messages, options=options)"
IO->>IO: _coerce_generation_options(options)
IO->>IO: _raise_on_unsupported_options(options, state)
IO->>IO: "records = []"
IO->>RM: are_tool_results_safe(messages)
RM->>RA: action.run()
RA->>CV: set(None)
RA->>ER: model_call / api_call
RA->>CV: set(RailLLMCall(usage, timing, model))
RA-->>RM: "RailResult(is_safe=True)"
RM->>CV: get_and_clear()
RM-->>IO: "RailResult(records=(RailCallRecord,))"
IO->>IO: records.extend(tool_result.records)
IO->>RM: is_input_safe(messages)
RM->>RA: action.run() [per input rail]
RA->>CV: set(RailLLMCall(...))
RM->>CV: get_and_clear()
RM-->>IO: "RailResult(records=(RailCallRecord,...))"
IO->>IO: records.extend(input_result.records)
IO->>ER: model_call(main, messages)
IO->>IO: records.append(_make_generation_record(...))
IO->>RM: is_output_safe(messages, response_text)
RM->>RA: action.run() [per output rail]
RM-->>IO: "RailResult(records=(RailCallRecord,...))"
IO->>IO: records.extend(output_result.records)
alt has_generation_options
IO->>IO: "log = _build_generation_log(records, options, total_duration)"
IO->>IO: _build_generation_response(text, reasoning, response, log)
IO-->>C: GenerationResponse
else bare path
IO-->>C: LLMMessage dict
end
%%{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 C as Caller
participant GR as Guardrails.generate_async
participant IO as IORails._do_generate
participant RM as RailsManager
participant RA as RailAction._get_llm_response
participant CV as _rail_llm_call_var (ContextVar)
participant ER as EngineRegistry
C->>GR: generate_async(messages, options)
GR->>IO: "_do_generate(messages, options=options)"
IO->>IO: _coerce_generation_options(options)
IO->>IO: _raise_on_unsupported_options(options, state)
IO->>IO: "records = []"
IO->>RM: are_tool_results_safe(messages)
RM->>RA: action.run()
RA->>CV: set(None)
RA->>ER: model_call / api_call
RA->>CV: set(RailLLMCall(usage, timing, model))
RA-->>RM: "RailResult(is_safe=True)"
RM->>CV: get_and_clear()
RM-->>IO: "RailResult(records=(RailCallRecord,))"
IO->>IO: records.extend(tool_result.records)
IO->>RM: is_input_safe(messages)
RM->>RA: action.run() [per input rail]
RA->>CV: set(RailLLMCall(...))
RM->>CV: get_and_clear()
RM-->>IO: "RailResult(records=(RailCallRecord,...))"
IO->>IO: records.extend(input_result.records)
IO->>ER: model_call(main, messages)
IO->>IO: records.append(_make_generation_record(...))
IO->>RM: is_output_safe(messages, response_text)
RM->>RA: action.run() [per output rail]
RM-->>IO: "RailResult(records=(RailCallRecord,...))"
IO->>IO: records.extend(output_result.records)
alt has_generation_options
IO->>IO: "log = _build_generation_log(records, options, total_duration)"
IO->>IO: _build_generation_response(text, reasoning, response, log)
IO-->>C: GenerationResponse
else bare path
IO-->>C: LLMMessage dict
end
Prompt To Fix All With AI
Fix the following 2 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 2
nemoguardrails/guardrails/actions/jailbreak_detection_action.py:44-47
The `return_value` here is a raw `bool` (the `jailbreak` field), while `content_safety_action` returns `{"allowed": bool, "policy_violations": list}` and `topic_safety_action` returns `{"on_topic": bool}`. Consumers iterating `GenerationLog.activated_rails[].executed_actions[].return_value` will see inconsistent types across rail kinds; a dict keeps all three uniform.
```suggestion
is_jailbreak = response["jailbreak"]
verdict = {"jailbreak": is_jailbreak, "score": score}
if is_jailbreak:
return RailResult(is_safe=False, reason=f"Score: {score}", return_value=verdict)
return RailResult(is_safe=True, reason=f"Score: {score}", return_value=verdict)
```
### Issue 2 of 2
nemoguardrails/guardrails/iorails.py:1060-1064
**Speculative path: main LLM duration absent from `llm_calls_duration`**
`_make_generation_record` is called with `None, None, None` for `started_at`, `finished_at`, and `duration`. In `_build_generation_stats`, `llm_calls_duration` is computed as `sum(r.duration or 0.0 for r in call_records)`, so the main LLM call contributes exactly `0.0` when speculative mode is active. A user who enables speculative generation and requests `options.log` will see `llm_calls_duration` that is always lower than the true value by the main inference time. `total_duration` is correct (it comes from `t_start`), so the discrepancy is visible to users comparing the two stats fields.
Reviews (2): Last reviewed commit: "Revert change to nemoguards config main ..." | Re-trigger Greptile
| is_jailbreak = response["jailbreak"] | ||
| if is_jailbreak: | ||
| return RailResult(is_safe=False, reason=f"Score: {score}", return_value=is_jailbreak) | ||
| return RailResult(is_safe=True, reason=f"Score: {score}", return_value=is_jailbreak) |
There was a problem hiding this comment.
return_value here is a raw bool (the jailbreak field), while content_safety_action returns {"allowed": bool, "policy_violations": list} and topic_safety_action returns {"on_topic": bool}. Consumers iterating GenerationLog.activated_rails[].executed_actions[].return_value will see inconsistent types across rail kinds; a dict keeps all three uniform.
| is_jailbreak = response["jailbreak"] | |
| if is_jailbreak: | |
| return RailResult(is_safe=False, reason=f"Score: {score}", return_value=is_jailbreak) | |
| return RailResult(is_safe=True, reason=f"Score: {score}", return_value=is_jailbreak) | |
| is_jailbreak = response["jailbreak"] | |
| verdict = {"jailbreak": is_jailbreak, "score": score} | |
| if is_jailbreak: | |
| return RailResult(is_safe=False, reason=f"Score: {score}", return_value=verdict) | |
| return RailResult(is_safe=True, reason=f"Score: {score}", return_value=verdict) |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/actions/jailbreak_detection_action.py
Line: 44-47
Comment:
The `return_value` here is a raw `bool` (the `jailbreak` field), while `content_safety_action` returns `{"allowed": bool, "policy_violations": list}` and `topic_safety_action` returns `{"on_topic": bool}`. Consumers iterating `GenerationLog.activated_rails[].executed_actions[].return_value` will see inconsistent types across rail kinds; a dict keeps all three uniform.
```suggestion
is_jailbreak = response["jailbreak"]
verdict = {"jailbreak": is_jailbreak, "score": score}
if is_jailbreak:
return RailResult(is_safe=False, reason=f"Score: {score}", return_value=verdict)
return RailResult(is_safe=True, reason=f"Score: {score}", return_value=verdict)
```
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!
| set_speculative_span_attrs(request_span, first_completed, "none") | ||
|
|
||
| log.debug("[%s] Main LLM response: %s", req_id, truncate(response.content)) | ||
| # Speculative races the main call, so per-call timing isn't tracked here; the | ||
| # generation record carries usage/model without timestamps or a duration. |
There was a problem hiding this comment.
llm_calls_duration
_make_generation_record is called with None, None, None for started_at, finished_at, and duration. In _build_generation_stats, llm_calls_duration is computed as sum(r.duration or 0.0 for r in call_records), so the main LLM call contributes exactly 0.0 when speculative mode is active. A user who enables speculative generation and requests options.log will see llm_calls_duration that is always lower than the true value by the main inference time. total_duration is correct (it comes from t_start), so the discrepancy is visible to users comparing the two stats fields.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/guardrails/iorails.py
Line: 1060-1064
Comment:
**Speculative path: main LLM duration absent from `llm_calls_duration`**
`_make_generation_record` is called with `None, None, None` for `started_at`, `finished_at`, and `duration`. In `_build_generation_stats`, `llm_calls_duration` is computed as `sum(r.duration or 0.0 for r in call_records)`, so the main LLM call contributes exactly `0.0` when speculative mode is active. A user who enables speculative generation and requests `options.log` will see `llm_calls_duration` that is always lower than the true value by the main inference time. `total_duration` is correct (it comes from `t_start`), so the discrepancy is visible to users comparing the two stats fields.
How can I resolve this? If you propose a fix, please make it concise.1 participant
Description
Adds the
GenerationResponsereturn argument to non-streaminggenerate()andgenerate_sync()methods. When aGenerationOptionsobject is provided, theGenerationResponsestructured response is returned rather than a dict withrole/content/tool_callskeys.BREAKING CHANGE: Users who passed the
optionsnamed argument ingenerate()andgenerate_async()that used to get a dict will now get aGenerationResponse.The
GenerationResponsecollects many different items of metadata, timestamps and durations into one place. TheGenerationResponse.logfield has a typeGenerationLog, which in IORails does not support theinternal_eventsandcolang_historyfields. TheGenerationLogOptionsclass controls which fields inGenerationLogare returned. IfGenerationLogOptions.internal_eventsorGenerationLogOptions.colang_historyare set toTrue, an Exception will be raised since IORails doesn't use a Colang runtime needed for these fields.Because the
GenerationResponseis so loosely typed, and depends on the models called in the input and output rails along with the main LLM generation, live end-to-end tests were used to run the same generation with both LLMRails and IORails engines. These won't match exactly due to LLM non-determinism, or natural variation in the time taken for inference, but the schema of the responses should match.GenerationResponsedetailed field support in IORailsFully supported
responsemessages, so always a 1-element[assistant_msg]list.tool_callsToolCall.to_dict()dict-args shape, matching LLMRails'GenerationResponse.tool_calls. (Bare/streaming paths keep the JSON-string wire shape, unchanged.)reasoning_contentreasoningfield, or<think>extracted from the completion; the structured path keepsresponsecontent clean (no inline<think>).llm_outputNone— identical to LLMRails, whose source for this field is never populated either.llm_output=Trueis accepted as a silent no-op.Partially supported
logactivated_rails(one syntheticExecutedAction+LLMCallInfoper rail), the flatllm_callslist, and aggregatestats— all synthesized from per-railRailCallRecords + the main call. Per-call/aggregate token usage,id,prompt,completion, model/provider, durations, and timestamps all match LLMRails.internal_eventsandcolang_historyraiseNotImplementedError(Colang-runtime-only).activated_rails[].decisionsis always[](a Colang construct). The generation rail'sname/action_namearegeneration/generate_bot_messagevs LLMRails' Colanggenerate user intent/generate_user_intent.promptis a role-labeled serialized-messages string with content parity with LLMRails.llm_metadataprovider_metadataverbatim (e.g.response_headers).usagesub-key — all token usage lives inlog(LLMRails mirror). OftenNonewhen the provider returns no metadata blob.Not supported
output_dataNone; passingoptions.output_varsraisesValueError.stateNone; passingoptions.stateraisesValueError.Integration-testing (using
compare_generate.py)This script issues the identical request to both IORails and LLMRails, with the
GenerationLogOptionsfields that IORails doesn't support disabled.It was run as below from the Guardrails uv environment
This generated the
llmrails_generation_response.jsonandiorails_generation_response.jsonfiles attached to the PR.iorails_generation_response.json
llmrails_generation_response.json
The Field-by-field Comparison output is below. For the full log file see here
generation_compare.log .
Related Issue(s)
Verification
Pre-commit
Unit-test
Integration test with Chat
AI Assistance
Checklist
Summary by CodeRabbit