Skip to content

draft(iorails): Support per-tool input and output rails - #2289

Draft
JashG wants to merge 3 commits into
NVIDIA-NeMo:developfrom
JashG:jgulabrai/per-tool-call-rails-iorails
Draft

draft(iorails): Support per-tool input and output rails#2289
JashG wants to merge 3 commits into
NVIDIA-NeMo:developfrom
JashG:jgulabrai/per-tool-call-rails-iorails

Conversation

@JashG

@JashG JashG commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds initial IORails support for applying LLM-based guardrails to specific tool calls and tool results.

per_tool lives under the existing rails.tool_output and rails.tool_input sections. These map to config.rails.tool_output.per_tool and config.rails.tool_input.per_tool on RailsConfig.

models:
  - type: main
    engine: nim
    model: meta/llama-3.1-70b-instruct

rails:
  tool_output:
    flows:
      - tool call validation
    per_tool:
      run_sql:
        - check tool call $check=blocked_tables
      export_data:
        - check tool call $check=external_export $model=main

  tool_input:
    flows:
      - tool result validation
    per_tool:
      run_sql:
        - check tool call $check=sensitive_result

prompts:
  - task: blocked_tables
    content: |-
      Block SQL that reads from protected tables.
      Tool call: {{ tool_call }}
      Respond with ALLOW or BLOCK on the final line.

  - task: external_export
    content: |-
      Block exports to external destinations.
      Tool call: {{ tool_call }}
      Respond with ALLOW or BLOCK on the final line.

  - task: sensitive_result
    content: |-
      Block tool results containing sensitive data.
      Tool result: {{ tool_result }}
      Respond with ALLOW or BLOCK on the final line.

What's included

  • Per-tool configuration for both tool boundaries
  • Prompt-based check tool call action
  • Sequential block-if-any execution
  • Composition with existing structural tool rails
  • Request-level enable/disable and filtering
  • Configuration validation
  • Tool-name recovery from tool_call_id
  • Fail-closed verdict handling
  • Focused unit and regression tests

Remaining work

  • Ensure flows can be run in parallel (i.e. Implement parallel: true)
  • Add tool_name to generation-log records
  • Refactor of PerToolCheckAction to use the standard RailAction pipeline so per-tool context can reuse shared tracing, model selection, call recording, and error handling instead of duplicating that logic in a custom run().
  • [Tests] Add top-level streaming and non-streaming integration tests
  • [Validation] Decide whether prompt and model references require load-time validation
  • [Validation] Define behavior when configuration forces fallback to LLMRails
  • [Docs] Add user-facing documentation and examples
  • Any additional changes requested in the RFC

@github-actions github-actions Bot added status: needs triage New issues that have not yet been reviewed or categorized. size: L labels Aug 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size: L status: needs triage New issues that have not yet been reviewed or categorized.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant