feat(guardrails): Add support for tool call input and output rails - #572
Open
JashG wants to merge 16 commits into
Open
feat(guardrails): Add support for tool call input and output rails#572JashG wants to merge 16 commits into
JashG wants to merge 16 commits into
Conversation
JashG
force-pushed
the
jgulabrai/agentic-guardrails-poc-0.23
branch
from
July 6, 2026 16:22
b28d1ed to
1a6e023
Compare
Contributor
|
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds tool-call guardrails for ChangesTool Call Guardrails Feature
SDK Loader Import Order
Sequence Diagram(s)sequenceDiagram
participant Client
participant GuardrailsMiddleware
participant LLMRails
participant ToolRailActions
participant Backend
Client->>GuardrailsMiddleware: request with tool_calls/tool results
GuardrailsMiddleware->>LLMRails: run tool_input rails
LLMRails->>ToolRailActions: check_tool_result_linkage
ToolRailActions-->>LLMRails: pass/block
GuardrailsMiddleware->>Backend: forward request
Backend-->>GuardrailsMiddleware: response with tool_calls
GuardrailsMiddleware->>LLMRails: run tool_output rails
LLMRails->>ToolRailActions: check_tool_allowlist / check_tool_arguments / check_tool_schema
ToolRailActions-->>LLMRails: pass/block
LLMRails-->>GuardrailsMiddleware: logs + guarded response
GuardrailsMiddleware-->>Client: response or refusal
Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
JashG
force-pushed
the
jgulabrai/agentic-guardrails-poc-0.23
branch
from
July 6, 2026 16:33
3d8113b to
d00f1f0
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (5)
plugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.py (1)
490-513: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the parsed tool flows instead of re-parsing on every build.
flows_co_path.read_text()+parse_colang_file(...)runs on everybuild()call even thoughflows.co's content never changes. Combine with a fullcopy.deepcopy(config)(line 493) for a config that may include KB/embeddings data — this is unnecessary per-build cost. Since onlyconfig.flowsis being replaced with a new list (not mutated in place), amodel_copy()(pydantic shallow copy) is enough, and the parsed flows can be memoized once at module import/first use.♻️ Suggested refactor
+_TOOL_RAIL_FLOWS: list[dict] | None = None + + +def _load_tool_rail_flows() -> list[dict]: + global _TOOL_RAIL_FLOWS + if _TOOL_RAIL_FLOWS is None: + from nemoguardrails.colang import parse_colang_file + + flows_co_path = Path(_TOOL_RAILS_PATH) / "flows.co" + parsed = parse_colang_file("flows.co", flows_co_path.read_text()) + flows = parsed.get("flows", []) + for flow in flows: + flow["is_subflow"] = True + flow["is_system_flow"] = True + _TOOL_RAIL_FLOWS = flows + return _TOOL_RAIL_FLOWS + + async def build(self, config: LibraryRailsConfig) -> LLMRails: ... - config = copy.deepcopy(config) - flows_co_path = Path(_TOOL_RAILS_PATH) / "flows.co" - flows_co_content = flows_co_path.read_text() - parsed = parse_colang_file("flows.co", flows_co_content) - tool_rail_flows = parsed.get("flows", []) - for flow in tool_rail_flows: - flow["is_subflow"] = True - flow["is_system_flow"] = True - config.flows = list(config.flows or []) + tool_rail_flows + config = config.model_copy() + config.flows = list(config.flows or []) + _load_tool_rail_flows()🤖 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 `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.py` around lines 490 - 513, The tool-flow injection in LLMRails build is doing unnecessary per-call work: it deep-copies the whole config and re-reads/re-parses flows.co every time. Update the build path to use a shallow pydantic copy of config instead of copy.deepcopy, and memoize the parsed flows from flows_co_path/parse_colang_file so they’re loaded once and reused. Keep the existing LLMRails construction and action registrations unchanged, and make sure config.flows still receives a fresh combined list.plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/README.md (1)
1-82: 📐 Maintainability & Code Quality | 🔵 TrivialDoc mixes Diataxis quadrants; missing prerequisites/Next Steps.
The README combines a reference table (Flow Surface), a how-to (Configuration), and explanation (Runtime Context, Failure Behavior) on one page, with no prerequisites section up top or a "Next Steps" section at the end.
As per coding guidelines, "Each documentation page should fit ONE Diataxis quadrant; do not mix tutorials with reference tables or how-tos with architecture explanations; use cross-links instead," and "Always list prerequisites at the top of documentation pages before other content" / "Include 'Next Steps' section at the end with cross-links to related documentation 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 `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/README.md` around lines 1 - 82, The README is mixing multiple Diataxis intents on one page, so split or refocus the content for the `Tool Rails` doc into a single quadrant and move the other material to linked pages. Add a prerequisites section at the top before `Flow Surface`, and add a `Next Steps` section at the end with cross-links to related docs; keep the current reference-style details like `check tool allowlist`, `check tool arguments`, `check tool schema`, and `check tool result linkage` but use links instead of combining reference, how-to, and explanation in one page.Source: Coding guidelines
plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.py (1)
210-277: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead code:
saw_tool_resultnever changes the outcome.Both branches of the final
if not saw_tool_result: return True/return Truereturn the same value, so the flag has no effect on behavior. Either this masks an intended stricter check (e.g. requiring at least one validated exchange) or it's leftover and can be removed.🤖 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 `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.py` around lines 210 - 277, The `check_tool_result_linkage` flow has dead logic because `saw_tool_result` does not affect the return value, so the final conditional is redundant. In `actions.py`, update `validate_exchange` and the surrounding `for message in messages` handling so the return path reflects the intended behavior: either remove `saw_tool_result` entirely if no stricter validation is needed, or use it to enforce the desired tool-result requirement before returning from `check_tool_result_linkage`. Keep the change localized to this function and preserve the existing validation checks for `tool_call_id` and `tool_calls`.agents/governed-optimization-poc/README.md (1)
1-108: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDoc structure violates Diataxis/guideline rules.
No Prerequisites section at all (should be at top, before "Files"), and the page mixes how-to content (Setup/Demo Commands, 29-90) with explanation content (How It Works, 91-108). No "Next Steps" section at the end.
As per coding guidelines, "Always list prerequisites at the top of documentation pages before other content", "Each documentation page should fit ONE Diataxis quadrant; do not mix tutorials with reference tables or how-tos with architecture explanations; use cross-links instead", and "Include 'Next Steps' section at the end with cross-links to related documentation 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 `@agents/governed-optimization-poc/README.md` around lines 1 - 108, The README mixes multiple Diataxis quadrants and is missing required documentation sections. Add a Prerequisites section at the top before Files, then separate the how-to content in Setup/Demo Commands from the explanation content in How It Works by moving one into a linked companion page or replacing it with cross-links. Also append a Next Steps section at the end with links to related documentation, and keep the main page focused on a single documentation purpose.Source: Coding guidelines
agents/agentic-guardrails-poc/README.md (1)
1-106: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDoc structure violates Diataxis/guideline rules.
Three violations: (1) Prerequisites (24-36) placed after "Files" section, not at the top. (2) Page mixes how-to content (Setup/Demo Commands, 37-83) with explanation content (How It Works, 84-106) in one page. (3) No "Next Steps" section at the end.
As per coding guidelines, "Always list prerequisites at the top of documentation pages before other content", "Each documentation page should fit ONE Diataxis quadrant; do not mix tutorials with reference tables or how-tos with architecture explanations; use cross-links instead", and "Include 'Next Steps' section at the end with cross-links to related documentation 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 `@agents/agentic-guardrails-poc/README.md` around lines 1 - 106, Rework the README structure to follow the documentation guidelines: move the Prerequisites section to the top before Files, then split the mixed how-to and explanation content into separate pages so this page stays in one Diataxis quadrant. Keep the setup/demo steps in the main guide and move the How It Works architecture details into a separate reference/explanation doc, with cross-links between them. Add a Next Steps section at the end of the remaining page linking to the related documentation.Source: Coding guidelines
🤖 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 `@agents/agentic-guardrails-poc/agentic-guardrails-poc.yml`:
- Around line 26-35: Template the workspace prefix in the llms.llm model_name
instead of hardcoding default/, because setup.sh creates the VirtualModel under
${WORKSPACE}/${VIRTUAL_MODEL_NAME}. Update agentic-guardrails-poc.yml so the
configuration is rendered from WORKSPACE (or otherwise injected before nemo
runs) and keep the llms.llm entry aligned with the workspace-specific model
path.
In `@agents/agentic-guardrails-poc/setup.sh`:
- Around line 7-18: The agent config is still using a hardcoded model workspace,
so overriding WORKSPACE in setup.sh does not propagate to the VirtualModel
lookup. Update the setup flow around WORKSPACE, VIRTUAL_MODEL_NAME, and
AGENT_CONFIG_PATH so the generated agent config uses the same workspace-prefixed
model name that is used when creating the VirtualModel. Ensure the
agentic-guardrails-poc.yml template gets the dynamic WORKSPACE value instead of
default/agentic-guardrails-poc-model, so both creation and lookup stay in sync.
In `@agents/deep-research-report.md`:
- Around line 5-198: The document contains unresolved citation artifacts that
appear as raw placeholder tokens instead of readable references throughout the
report. Update the markdown in deep-research-report to remove these
citation-tool leftovers or convert them into proper markdown footnotes/links in
every affected section, including the intro, tables, workflows, checklist, and
references. Use the repeated inline citation symbols and citation clusters as
the locations to clean up, and ensure the final text reads cleanly without any
raw placeholder strings.
In `@agents/governed-optimization-poc/governed-optimization-poc.yml`:
- Around line 11-18: The llm configuration is hardcoding model_name instead of
deriving it from WORKSPACE, so non-default environments resolve to the wrong
VirtualModel. Update the governed-optimization-poc YAML to make the
llm.model_name value use the same WORKSPACE-based naming convention that
setup.sh expects, and keep the change localized to the llm configuration so the
model selection stays aligned with the active workspace.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.py`:
- Around line 878-889: The _reset method in llmrails_cache currently skips
update_llm() when main_llm is None, leaving stale llm_generation_actions.llm and
the "llm" action param behind in reused LLMRails instances. Update the None path
in _reset to clear all model-related references consistently, either by routing
through update_llm() or by explicitly resetting rails.llm,
rails.llm_generation_actions.llm, and the related action param together.
- Around line 463-470: The built wheel is missing the plugin’s bundled Colang
asset, so `LLMRailsCache` can fail when `flows_co_path.read_text()` looks for
`tool_rails/flows.co`. Update the packaging config so Hatchling includes
`src/nemo_guardrails_plugin/tool_rails/flows.co` in the wheel, using an
`include` or `force-include` rule, and verify the asset is shipped alongside the
`LLMRailsCache`/`_TOOL_RAILS_PATH` logic.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/responses.py`:
- Around line 130-138: The response parsing in the message extraction path can
still crash when choices[0]["message"] is null because the existing try/except
only catches lookup errors, not a None value. Update the logic in responses.py
around the assistant message construction so that after reading
response_result["choices"][0]["message"], you verify it is a dict before calling
message.get, and fall back to assistant_message when it is missing or invalid.
Keep the fix localized to the response handling helper that builds
assistant_message and preserves tool_calls.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/streaming.py`:
- Around line 288-295: The streamed tool-call id fallback in
_finalize_streaming_tool_call is content-based, so identical tool calls can
reuse the same id and break check_tool_result_linkage. Update the id generation
in _finalize_streaming_tool_call to use a per-call unique value from the stream
context, such as the tool-call index or another incrementing unique token, while
keeping the existing defaults for type and function fields.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.py`:
- Around line 146-148: Add the missing jsonschema dependency to the
nemo-guardrails plugin so check_tool_schema does not rely on a transitive
install. Update the plugin’s dependency declaration in the nemo-guardrails
pyproject so the import used in actions.py is guaranteed at runtime, and keep
the existing import inside check_tool_schema aligned with that declared
dependency.
---
Nitpick comments:
In `@agents/agentic-guardrails-poc/README.md`:
- Around line 1-106: Rework the README structure to follow the documentation
guidelines: move the Prerequisites section to the top before Files, then split
the mixed how-to and explanation content into separate pages so this page stays
in one Diataxis quadrant. Keep the setup/demo steps in the main guide and move
the How It Works architecture details into a separate reference/explanation doc,
with cross-links between them. Add a Next Steps section at the end of the
remaining page linking to the related documentation.
In `@agents/governed-optimization-poc/README.md`:
- Around line 1-108: The README mixes multiple Diataxis quadrants and is missing
required documentation sections. Add a Prerequisites section at the top before
Files, then separate the how-to content in Setup/Demo Commands from the
explanation content in How It Works by moving one into a linked companion page
or replacing it with cross-links. Also append a Next Steps section at the end
with links to related documentation, and keep the main page focused on a single
documentation purpose.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.py`:
- Around line 490-513: The tool-flow injection in LLMRails build is doing
unnecessary per-call work: it deep-copies the whole config and
re-reads/re-parses flows.co every time. Update the build path to use a shallow
pydantic copy of config instead of copy.deepcopy, and memoize the parsed flows
from flows_co_path/parse_colang_file so they’re loaded once and reused. Keep the
existing LLMRails construction and action registrations unchanged, and make sure
config.flows still receives a fresh combined list.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.py`:
- Around line 210-277: The `check_tool_result_linkage` flow has dead logic
because `saw_tool_result` does not affect the return value, so the final
conditional is redundant. In `actions.py`, update `validate_exchange` and the
surrounding `for message in messages` handling so the return path reflects the
intended behavior: either remove `saw_tool_result` entirely if no stricter
validation is needed, or use it to enforce the desired tool-result requirement
before returning from `check_tool_result_linkage`. Keep the change localized to
this function and preserve the existing validation checks for `tool_call_id` and
`tool_calls`.
In `@plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/README.md`:
- Around line 1-82: The README is mixing multiple Diataxis intents on one page,
so split or refocus the content for the `Tool Rails` doc into a single quadrant
and move the other material to linked pages. Add a prerequisites section at the
top before `Flow Surface`, and add a `Next Steps` section at the end with
cross-links to related docs; keep the current reference-style details like
`check tool allowlist`, `check tool arguments`, `check tool schema`, and `check
tool result linkage` but use links instead of combining reference, how-to, and
explanation in one page.
🪄 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: b6658715-087d-4f03-bbe2-6c6b5ce5f5d6
📒 Files selected for processing (30)
agents/agentic-guardrails-linear-issue.mdagents/agentic-guardrails-poc/README.mdagents/agentic-guardrails-poc/agentic-guardrails-poc.ymlagents/agentic-guardrails-poc/setup.shagents/agentic-guardrails-poc/tool-rails-config.jsonagents/agentic-guardrails-proposal.mdagents/deep-research-report.mdagents/governed-optimization-poc/README.mdagents/governed-optimization-poc/governed-optimization-poc.ymlagents/governed-optimization-poc/pyproject.tomlagents/governed-optimization-poc/setup.shagents/governed-optimization-poc/src/governed_optimization_agent/register.pyagents/governed-optimization-poc/tool-rails-config.jsonplugins/nemo-guardrails/src/nemo_guardrails_plugin/constants.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/middleware.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/rails.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/responses.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/streaming.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/README.mdplugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/flows.coplugins/nemo-guardrails/tests/integration/test_llmrails_cache_real_build.pyplugins/nemo-guardrails/tests/integration/test_tool_rails.pyplugins/nemo-guardrails/tests/unit/test_llmrails_cache.pyplugins/nemo-guardrails/tests/unit/test_middleware.pyplugins/nemo-guardrails/tests/unit/test_rails.pyplugins/nemo-guardrails/tests/unit/test_responses.pyplugins/nemo-guardrails/tests/unit/test_streaming.pyplugins/nemo-guardrails/tests/unit/test_tool_rails.py
Comment on lines
+7
to
+18
| WORKSPACE="${WORKSPACE:-default}" | ||
| CONFIG_NAME="${CONFIG_NAME:-agentic-tool-rails-poc}" | ||
| VIRTUAL_MODEL_NAME="${VIRTUAL_MODEL_NAME:-agentic-guardrails-poc-model}" | ||
| AGENT_NAME="${AGENT_NAME:-agentic-guardrails-poc}" | ||
| BACKEND_MODEL="${BACKEND_MODEL:-default/nvidia-nemotron-3-nano-30b-a3b}" | ||
| DEPLOY_AGENT="${DEPLOY_AGENT:-0}" | ||
| RECREATE_AGENT="${RECREATE_AGENT:-1}" | ||
| RECREATE_VIRTUAL_MODEL="${RECREATE_VIRTUAL_MODEL:-1}" | ||
|
|
||
| CONFIG_PATH="${SCRIPT_DIR}/tool-rails-config.json" | ||
| AGENT_CONFIG_PATH="${SCRIPT_DIR}/agentic-guardrails-poc.yml" | ||
| GUARDRAIL_CONFIG_ID="${WORKSPACE}/${CONFIG_NAME}" |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
WORKSPACE override not reflected in agent config's model_name.
WORKSPACE is used here to create the VirtualModel (Line 62-66), but agentic-guardrails-poc.yml's model_name hardcodes default/agentic-guardrails-poc-model (see review on that file). If a user overrides WORKSPACE, the agent will look up the model in the wrong workspace and setup will silently produce a broken demo.
Also applies to: 62-66
🤖 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 `@agents/agentic-guardrails-poc/setup.sh` around lines 7 - 18, The agent config
is still using a hardcoded model workspace, so overriding WORKSPACE in setup.sh
does not propagate to the VirtualModel lookup. Update the setup flow around
WORKSPACE, VIRTUAL_MODEL_NAME, and AGENT_CONFIG_PATH so the generated agent
config uses the same workspace-prefixed model name that is used when creating
the VirtualModel. Ensure the agentic-guardrails-poc.yml template gets the
dynamic WORKSPACE value instead of default/agentic-guardrails-poc-model, so both
creation and lookup stay in sync.
Comment on lines
+11
to
+18
| llms: | ||
| llm: | ||
| _type: openai | ||
| api_key: not-used | ||
| model_name: default/governed-optimization-poc-model | ||
| temperature: 0.0 | ||
| max_tokens: 1024 | ||
| disable_streaming: true |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== files ==\n'
git ls-files 'agents/governed-optimization-poc/*' | sed 's#^`#-` #'
printf '\n== outline governed-optimization-poc.yml ==\n'
ast-grep outline agents/governed-optimization-poc/governed-optimization-poc.yml --view expanded || true
printf '\n== relevant text search ==\n'
rg -n "WORKSPACE|default/governed-optimization-poc-model|model_name|workspace" agents/governed-optimization-poc -S || true
printf '\n== README and setup.sh snippets ==\n'
for f in agents/governed-optimization-poc/README.md agents/governed-optimization-poc/setup.sh; do
if [ -f "$f" ]; then
echo "--- $f (first 220 lines) ---"
nl -ba "$f" | sed -n '1,220p'
fi
doneRepository: NVIDIA-NeMo/nemo-platform
Length of output: 4067
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- README.md ---'
sed -n '1,160p' agents/governed-optimization-poc/README.md
echo '--- setup.sh ---'
sed -n '1,180p' agents/governed-optimization-poc/setup.sh
echo '--- governed-optimization-poc.yml ---'
sed -n '1,120p' agents/governed-optimization-poc/governed-optimization-poc.ymlRepository: NVIDIA-NeMo/nemo-platform
Length of output: 10502
Tie model_name to WORKSPACE setup.sh already honors WORKSPACE, but this hardcoded default/governed-optimization-poc-model makes non-default setups point at the wrong VirtualModel.
🤖 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 `@agents/governed-optimization-poc/governed-optimization-poc.yml` around lines
11 - 18, The llm configuration is hardcoding model_name instead of deriving it
from WORKSPACE, so non-default environments resolve to the wrong VirtualModel.
Update the governed-optimization-poc YAML to make the llm.model_name value use
the same WORKSPACE-based naming convention that setup.sh expects, and keep the
change localized to the llm configuration so the model selection stays aligned
with the active workspace.
JashG
force-pushed
the
jgulabrai/agentic-guardrails-poc-0.23
branch
from
July 6, 2026 17:55
d00f1f0 to
54a777c
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
agents/governed-optimization-poc/src/governed_optimization_agent/register.py (1)
73-87: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDedupe the repetitive
add_functionregistration.Same
if name in selected: group.add_function(...)pattern repeated five times. A dict of name→fn plus a loop is more concise and scales better as tools are added.♻️ Proposed refactor
- selected = set(config.include) - if "search_internal_knowledge" in selected: - group.add_function( - name="search_internal_knowledge", - fn=_search_internal_knowledge, - description=_search_internal_knowledge.__doc__, - ) - if "estimate_cost" in selected: - group.add_function(name="estimate_cost", fn=_estimate_cost, description=_estimate_cost.__doc__) - if "run_eval" in selected: - group.add_function(name="run_eval", fn=_run_eval, description=_run_eval.__doc__) - if "propose_update" in selected: - group.add_function(name="propose_update", fn=_propose_update, description=_propose_update.__doc__) - if "deploy_candidate" in selected: - group.add_function(name="deploy_candidate", fn=_deploy_candidate, description=_deploy_candidate.__doc__) + available = { + "search_internal_knowledge": _search_internal_knowledge, + "estimate_cost": _estimate_cost, + "run_eval": _run_eval, + "propose_update": _propose_update, + "deploy_candidate": _deploy_candidate, + } + for name in config.include: + fn = available.get(name) + if fn is not None: + group.add_function(name=name, fn=fn, description=fn.__doc__)🤖 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 `@agents/governed-optimization-poc/src/governed_optimization_agent/register.py` around lines 73 - 87, The function registration in register.py is repeating the same selected-membership check and group.add_function call for each tool, so refactor this block into a name-to-function mapping and iterate over it. Keep the existing behavior in the registration logic around group.add_function, _search_internal_knowledge, _estimate_cost, _run_eval, _propose_update, and _deploy_candidate, but drive it from a single loop so adding new tools only requires updating the mapping.
🤖 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.
Nitpick comments:
In
`@agents/governed-optimization-poc/src/governed_optimization_agent/register.py`:
- Around line 73-87: The function registration in register.py is repeating the
same selected-membership check and group.add_function call for each tool, so
refactor this block into a name-to-function mapping and iterate over it. Keep
the existing behavior in the registration logic around group.add_function,
_search_internal_knowledge, _estimate_cost, _run_eval, _propose_update, and
_deploy_candidate, but drive it from a single loop so adding new tools only
requires updating the mapping.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 54b722cd-8b9c-466b-81ef-c616dcb2ff23
📒 Files selected for processing (30)
agents/agentic-guardrails-linear-issue.mdagents/agentic-guardrails-poc/README.mdagents/agentic-guardrails-poc/agentic-guardrails-poc.ymlagents/agentic-guardrails-poc/setup.shagents/agentic-guardrails-poc/tool-rails-config.jsonagents/agentic-guardrails-proposal.mdagents/deep-research-report.mdagents/governed-optimization-poc/README.mdagents/governed-optimization-poc/governed-optimization-poc.ymlagents/governed-optimization-poc/pyproject.tomlagents/governed-optimization-poc/setup.shagents/governed-optimization-poc/src/governed_optimization_agent/register.pyagents/governed-optimization-poc/tool-rails-config.jsonplugins/nemo-guardrails/src/nemo_guardrails_plugin/constants.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/middleware.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/rails.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/responses.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/streaming.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/README.mdplugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.pyplugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/flows.coplugins/nemo-guardrails/tests/integration/test_llmrails_cache_real_build.pyplugins/nemo-guardrails/tests/integration/test_tool_rails.pyplugins/nemo-guardrails/tests/unit/test_llmrails_cache.pyplugins/nemo-guardrails/tests/unit/test_middleware.pyplugins/nemo-guardrails/tests/unit/test_rails.pyplugins/nemo-guardrails/tests/unit/test_responses.pyplugins/nemo-guardrails/tests/unit/test_streaming.pyplugins/nemo-guardrails/tests/unit/test_tool_rails.py
✅ Files skipped from review due to trivial changes (5)
- agents/governed-optimization-poc/pyproject.toml
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/README.md
- agents/governed-optimization-poc/README.md
- plugins/nemo-guardrails/tests/unit/test_llmrails_cache.py
- agents/agentic-guardrails-poc/README.md
🚧 Files skipped from review as they are similar to previous changes (18)
- agents/governed-optimization-poc/tool-rails-config.json
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/constants.py
- agents/governed-optimization-poc/governed-optimization-poc.yml
- agents/agentic-guardrails-poc/agentic-guardrails-poc.yml
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/flows.co
- agents/agentic-guardrails-poc/tool-rails-config.json
- plugins/nemo-guardrails/tests/unit/test_streaming.py
- plugins/nemo-guardrails/tests/integration/test_tool_rails.py
- agents/agentic-guardrails-poc/setup.sh
- plugins/nemo-guardrails/tests/unit/test_responses.py
- agents/governed-optimization-poc/setup.sh
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/tool_rails/actions.py
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/rails.py
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/llmrails_cache.py
- plugins/nemo-guardrails/tests/integration/test_llmrails_cache_real_build.py
- plugins/nemo-guardrails/src/nemo_guardrails_plugin/middleware.py
- plugins/nemo-guardrails/tests/unit/test_rails.py
- plugins/nemo-guardrails/tests/unit/test_middleware.py
JashG
force-pushed
the
jgulabrai/agentic-guardrails-poc-0.23
branch
from
July 6, 2026 17:59
bf5522b to
18d702c
Compare
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
…utput rails Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
JashG
force-pushed
the
jgulabrai/agentic-guardrails-poc-0.23
branch
from
July 7, 2026 18:56
d37086f to
7bb2032
Compare
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
JashG
force-pushed
the
jgulabrai/agentic-guardrails-poc-0.23
branch
from
July 14, 2026 15:36
10396c5 to
55d88e2
Compare
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit