Skip to content

Conversation

@thepatrickchin
Copy link
Contributor

@thepatrickchin thepatrickchin commented Sep 9, 2025

Description

  • Added optional return_direct: list[str] parameter to constructor, allowing specified tools to return their results directly to the user without additional LLM processing.
  • Implemented tool_conditional_edge() method to determine when tools should return directly
  • Builds modified state graph following the established pattern from ReWOO agent

Closes #758

By Submitting this PR I confirm:

  • I am familiar with the Contributing Guidelines.
  • We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
    • Any contribution which contains commits that are not Signed-Off will not be accepted.
  • When the PR is ready for review, new or existing tests cover these changes.
  • When the PR is ready for review, the documentation is up to date with these changes.

Summary by CodeRabbit

  • New Features

    • Added an optional return_direct setting to return outputs from specified tools directly to users, enabling early termination of the tool path and faster responses.
  • Documentation

    • Workflow docs updated to describe the new return_direct option alongside existing configuration settings.
  • Tests

    • Expanded tests to cover return_direct routing, match/no-match edge cases, and compatibility with existing flows.

@copy-pr-bot
Copy link

copy-pr-bot bot commented Sep 9, 2025

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai
Copy link

coderabbitai bot commented Sep 9, 2025

Walkthrough

Adds a configurable return_direct option to the tool-calling agent; implements conditional graph edges and decision logic to allow specified tool responses to be returned directly, updates workflow config/registration and docs, and extends tests for the new behavior.

Changes

Cohort / File(s) Summary
Documentation: return_direct option
docs/source/workflows/about/tool-calling-agent.md
Adds return_direct as an optional list of tool names in Configurable Options describing direct-return behavior (bypass agent reasoning).
Agent graph: conditional end after tools
src/nat/agent/tool_calling_agent/agent.py
Adds return_direct parameter to ToolCallAgentGraph.__init__ and stores tool names. Adds async def tool_conditional_edge(...) to decide END vs TOOL based on the last ToolMessage. Introduces _build_graph to assemble a StateGraph/CompiledStateGraph with conditional edges when return_direct is set; updates build_graph to delegate and adds imports for ToolMessage, StateGraph, and CompiledStateGraph.
Workflow config & registration
src/nat/agent/tool_calling_agent/register.py
Adds `return_direct: list[FunctionRef]
Tests: return_direct behavior
tests/nat/agent/test_tool_calling.py
Adds fixtures and tests for return_direct mode covering graph structure, conditional edges, branch ends, and final messages when return_direct is active; updates/extends tests for non-return_direct behavior and edge cases (missing name, empty messages).

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor U as User
  participant G as ToolCallAgentGraph
  participant L as Agent LLM
  participant T as Tool(s)

  U->>G: Input message
  G->>L: Agent LLM decides next action
  alt L decides to call tool
    G->>T: Invoke tool
    T-->>G: ToolMessage(name, output)
    alt name in return_direct
      rect rgb(230,248,230)
        note right of G: Conditional edge -> END (direct return)
        G-->>U: Return ToolMessage output directly
      end
    else name not in return_direct
      rect rgb(232,240,255)
        G->>L: Continue LLM processing with tool output
        L-->>G: Final response
        G-->>U: Return LLM-crafted answer
      end
    end
  else No tool call
    L-->>G: Final response
    G-->>U: Return LLM-crafted answer
  end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 3.85% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check ✅ Passed The title succinctly and imperatively describes the main change—adding a return_direct option to the tool_calling_agent to return tool outputs directly—and it matches the PR's code, tests, and documentation changes while staying within the ~72-character guideline.
Linked Issues Check ✅ Passed The PR implements the core objective from issue [#758] by introducing a return_direct configuration, adding tool_conditional_edge logic and modified graph construction to early-exit when a listed tool returns a ToolMessage, exposing return_direct in the workflow config with FunctionRef→tool conversion, and adding tests and docs to cover the behavior, which collectively satisfy the coding requirement to bypass LLM processing for specified tools. The changes focus on the tool-calling paths and include unit tests and documentation updates, so the linked-issue requirements appear met. Please verify the build_graph signature consistency and the runtime mapping of configured FunctionRef names to actual tool.name values as a final integration check.
Out of Scope Changes Check ✅ Passed All modified files and changes are directly tied to implementing return_direct (agent logic, graph construction, workflow config, tests, and docs) and I do not see edits touching unrelated subsystems or features, so the PR appears free of out-of-scope changes; the only minor issue to confirm is an internal inconsistency in summaries regarding build_graph parameter handling, which is an implementation detail rather than an unrelated change.
✨ Finishing touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Tip

👮 Agentic pre-merge checks are now available in preview!

Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.

  • Built-in checks – Quickly apply ready-made checks to enforce title conventions, require pull request descriptions that follow templates, validate linked issues for compliance, and more.
  • Custom agentic checks – Define your own rules using CodeRabbit’s advanced agentic capabilities to enforce organization-specific policies and workflows. For example, you can instruct CodeRabbit’s agent to verify that API documentation is updated whenever API schema files are modified in a PR. Note: Upto 5 custom checks are currently allowed during the preview period. Pricing for this feature will be announced in a few weeks.

Please see the documentation for more information.

Example:

reviews:
  pre_merge_checks:
    custom_checks:
      - name: "Undocumented Breaking Changes"
        mode: "warning"
        instructions: |
          Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).

Please share your feedback with us on this Discord post.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai bot added feature request New feature or request non-breaking Non-breaking change labels Sep 9, 2025
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
src/nat/agent/tool_calling_agent/agent.py (2)

19-27: Add missing imports for stronger typing and safer message checks.
Import CompiledGraph for return types and ToolMessage for explicit isinstance checks.

 from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import SystemMessage
+from langchain_core.messages import SystemMessage, ToolMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.runnables import RunnableLambda
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
+from langgraph.graph.graph import CompiledGraph

52-62: Annotate init and normalize/validate return_direct against provided tools.

  • Public API requires type hints; add -> None.
  • Normalize to a frozenset, warn on unknown tool names; avoids subtle mismatches and speeds membership checks.
-    def __init__(
+    def __init__(
         self,
         llm: BaseChatModel,
         tools: list[BaseTool],
         prompt: str | None = None,
         callbacks: list[AsyncCallbackHandler] = None,
         detailed_logs: bool = False,
         log_response_max_chars: int = 1000,
         handle_tool_errors: bool = True,
-        return_direct: list[str] | None = None,
-    ):
+        return_direct: list[str] | None = None,
+    ) -> None:
@@
-        self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
-        self.return_direct = return_direct
+        self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
+        # Normalize and validate return_direct
+        if return_direct:
+            allowed_tools = {t.name for t in tools}
+            normalized = {n.strip() for n in return_direct if n and isinstance(n, str)}
+            unknown = normalized - allowed_tools
+            if unknown:
+                logger.warning("%s Ignoring unknown tools in return_direct: %s", AGENT_LOG_PREFIX, sorted(unknown))
+            self.return_direct = frozenset(normalized & allowed_tools)
+        else:
+            self.return_direct = frozenset()
🧹 Nitpick comments (11)
src/nat/agent/tool_calling_agent/agent.py (1)

152-169: Behavioral caveat: direct-return on tool errors.
If ToolNode returns an error ToolMessage (handle_tool_errors=True), this path will still return directly. Consider documenting this or adding a heuristic/flag to route errors back through the agent for retries.

Would you like a follow-up patch adding an opt-in return_direct_on_error: bool = False and a basic error heuristic?

docs/source/workflows/about/tool-calling-agent.md (2)

77-91: Unify list markers and clarify return_direct note.
Switch bullets to dashes for markdownlint (MD004) and add a brief caveat.

-### Configurable Options
+### Configurable Options
@@
-* `tool_names`: A list of tools that the agent can call. The tools must be functions configured in the YAML file
+ - `tool_names`: A list of tools that the agent can call. The tools must be functions configured in the YAML file
@@
-* `llm_name`: The LLM the agent should use. The LLM must be configured in the YAML file
+ - `llm_name`: The LLM the agent should use. The LLM must be configured in the YAML file
@@
-* `verbose`: Defaults to False (useful to prevent logging of sensitive data). If set to True, the agent will log input, output, and intermediate steps.
+ - `verbose`: Defaults to False (useful to prevent logging of sensitive data). If set to True, the agent will log input, output, and intermediate steps.
@@
-* `handle_tool_errors`: Defaults to True. All tool errors will be caught and a `ToolMessage` with an error message will be returned, allowing the agent to retry.
+ - `handle_tool_errors`: Defaults to True. All tool errors will be caught and a `ToolMessage` with an error message will be returned, allowing the agent to retry.
@@
-* `max_iterations`: Defaults to 15. The maximum number of tool calls the agent may perform.
+ - `max_iterations`: Defaults to 15. The maximum number of tool calls the agent may perform.
@@
-* `return_direct`: Optional list of tool names that should return their output directly without additional agent processing. When a tool in this list is called, its response is returned immediately to the user, bypassing the agent's reasoning step.
+ - `return_direct`: Optional list of tool names that should return their output directly without additional agent processing. When a tool in this list is called, its response is returned immediately to the user, bypassing the agent's reasoning step. Note: This is only supported by the tool-calling agent and may conflict with agents that rely on intermediate reasoning (e.g., ReAct).
@@
-* `description`:  Defaults to "Tool Calling Agent Workflow". When the agent is configured as a function, this config option allows us to control the tool description (for example, when used as a tool within another agent).
+ - `description`:  Defaults to "Tool Calling Agent Workflow". When the agent is configured as a function, this config option allows us to control the tool description (for example, when used as a tool within another agent).

77-91: Consider adding a YAML example showing return_direct usage.
A short snippet will make the option discoverable and testable.

Would you like me to propose a doc patch with an example?

src/nat/agent/tool_calling_agent/register.py (2)

50-52: Config field looks good; ensure naming expectations are clear.
If matching is case-sensitive, consider documenting it. Otherwise, normalize to lower-case in code/doc.


73-81: Validate return_direct against actual tool names to avoid silent misconfig.
Warn and drop unknown entries before constructing the graph.

@@
-    # construct the Tool Calling Agent Graph from the configured llm, and tools
-    graph: CompiledGraph = await ToolCallAgentGraph(llm=llm,
+    # Validate return_direct against available tools
+    rd = config.return_direct
+    if rd:
+        tool_names = {t.name for t in tools}
+        normalized = {n.strip() for n in rd if n and isinstance(n, str)}
+        unknown = normalized - tool_names
+        if unknown:
+            logger.warning("%s Ignoring unknown tools in return_direct: %s", AGENT_LOG_PREFIX, sorted(unknown))
+        rd = sorted(normalized & tool_names)
+
+    # construct the Tool Calling Agent Graph from the configured llm, and tools
+    graph: CompiledGraph = await ToolCallAgentGraph(llm=llm,
                                                     tools=tools,
                                                     prompt=prompt,
                                                     detailed_logs=config.verbose,
                                                     log_response_max_chars=config.log_response_max_chars,
-                                                    handle_tool_errors=config.handle_tool_errors,
-                                                    return_direct=config.return_direct).build_graph()
+                                                    handle_tool_errors=config.handle_tool_errors,
+                                                    return_direct=rd).build_graph()
tests/nat/agent/test_tool_calling.py (6)

101-115: Make assertions resilient to ordering; compare sets instead of lists.

Dict key order and tool collections aren’t guaranteed. Comparing sets avoids brittle tests.

-    assert list(agent.tool_caller.tools_by_name.keys()) == ['Tool A', 'Tool B']
-    assert agent.return_direct == return_direct_tools
+    assert set(agent.tool_caller.tools_by_name.keys()) == {'Tool A', 'Tool B'}
+    assert set(agent.return_direct) == set(return_direct_tools)

101-115: Silence Ruff S101 for pytest-style asserts in tests.

Ruff flags S101 here and elsewhere in this file. If your pyproject doesn’t already ignore S101 for tests, add a file-level directive.

Add below the SPDX header lines:

# ruff: noqa: S101  # pytest-style asserts are expected in tests

Alternatively, configure a per-directory ignore for tests in pyproject.toml.


124-131: Deduplicate return_direct tool names across tests.

Define a module-level constant to avoid divergence.

RETURN_DIRECT_TOOLS = ['Tool A']  # at module top

# then reuse
agent = ToolCallAgentGraph(..., return_direct=RETURN_DIRECT_TOOLS)

144-156: Avoid node-order brittleness; assert on sets.

Graph node ordering isn’t a public contract. Use set equality for stability.

-    assert list(graph.nodes.keys()) == ['__start__', 'agent', 'tool']
+    assert set(graph.nodes.keys()) == {'__start__', 'agent', 'tool'}

268-275: Strengthen assertion to verify raw tool output is surfaced.

Assert the ToolMessage content to confirm the direct-return behavior passes through the tool’s raw payload.

-    assert isinstance(last_message, ToolMessage)
-    assert last_message.name == 'Tool A'
+    assert isinstance(last_message, ToolMessage)
+    assert last_message.name == 'Tool A'
+    assert last_message.content == 'mock query'

189-221: Add test for ToolMessage without name attribute
Add a test where the final message is a ToolMessage constructed without name to confirm tool_conditional_edge falls back to AgentDecision.TOOL.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between f5836c0 and 16c3762.

📒 Files selected for processing (4)
  • docs/source/workflows/about/tool-calling-agent.md (1 hunks)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
  • src/nat/agent/tool_calling_agent/register.py (2 hunks)
  • tests/nat/agent/test_tool_calling.py (4 hunks)
🧰 Additional context used
📓 Path-based instructions (10)
docs/source/**/*.md

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

docs/source/**/*.md: Use the official naming: first use “NVIDIA NeMo Agent toolkit”; subsequent uses “NeMo Agent toolkit”; never use deprecated names in documentation
Documentation sources must be Markdown under docs/source; keep docs in sync and fix Sphinx errors/broken links
Documentation must be clear, comprehensive, free of TODO/FIXME/placeholders/offensive/outdated terms; fix spelling; adhere to Vale vocab allow/reject lists

Files:

  • docs/source/workflows/about/tool-calling-agent.md
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • docs/source/workflows/about/tool-calling-agent.md
  • src/nat/agent/tool_calling_agent/register.py
  • tests/nat/agent/test_tool_calling.py
  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • docs/source/workflows/about/tool-calling-agent.md
  • src/nat/agent/tool_calling_agent/register.py
  • tests/nat/agent/test_tool_calling.py
  • src/nat/agent/tool_calling_agent/agent.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • docs/source/workflows/about/tool-calling-agent.md
  • src/nat/agent/tool_calling_agent/register.py
  • tests/nat/agent/test_tool_calling.py
  • src/nat/agent/tool_calling_agent/agent.py
docs/source/**/*

⚙️ CodeRabbit configuration file

This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the docs/source/_static directory.

Files:

  • docs/source/workflows/about/tool-calling-agent.md
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/register.py
  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/register.py
  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/register.py
  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

Files:

  • src/nat/agent/tool_calling_agent/register.py
  • tests/nat/agent/test_tool_calling.py
  • src/nat/agent/tool_calling_agent/agent.py
tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Unit tests must live under tests/ and use configured markers (e2e, integration, etc.)

Files:

  • tests/nat/agent/test_tool_calling.py

⚙️ CodeRabbit configuration file

tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using the test_ prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using the fixture_ prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass

Files:

  • tests/nat/agent/test_tool_calling.py
**/tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/tests/**/*.py: Test functions must use the test_ prefix and snake_case
Extract repeated test code into pytest fixtures; fixtures should set name=... in @pytest.fixture and functions named with fixture_ prefix
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Use pytest with pytest-asyncio for async code; mock external services with pytest_httpserver or unittest.mock

Files:

  • tests/nat/agent/test_tool_calling.py
🧬 Code graph analysis (3)
src/nat/agent/tool_calling_agent/register.py (1)
src/nat/agent/tool_calling_agent/agent.py (1)
  • build_graph (197-204)
tests/nat/agent/test_tool_calling.py (3)
tests/conftest.py (2)
  • mock_llm (274-332)
  • mock_tool (336-358)
src/nat/agent/base.py (1)
  • AgentDecision (62-64)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
src/nat/agent/tool_calling_agent/agent.py (3)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/rewoo_agent/agent.py (3)
  • _build_graph (344-367)
  • conditional_edge (327-342)
  • build_graph (369-376)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
🪛 markdownlint-cli2 (0.17.2)
docs/source/workflows/about/tool-calling-agent.md

89-89: Unordered list style
Expected: dash; Actual: asterisk

(MD004, ul-style)

🪛 Ruff (0.12.2)
tests/nat/agent/test_tool_calling.py

108-108: Use of assert detected

(S101)


109-109: Use of assert detected

(S101)


110-110: Use of assert detected

(S101)


111-111: Use of assert detected

(S101)


112-112: Use of assert detected

(S101)


113-113: Use of assert detected

(S101)


114-114: Use of assert detected

(S101)


146-146: Use of assert detected

(S101)


147-147: Use of assert detected

(S101)


148-148: Use of assert detected

(S101)


149-149: Use of assert detected

(S101)


153-153: Use of assert detected

(S101)


154-154: Use of assert detected

(S101)


155-155: Use of assert detected

(S101)


193-193: Use of assert detected

(S101)


200-200: Use of assert detected

(S101)


207-207: Use of assert detected

(S101)


214-214: Use of assert detected

(S101)


220-220: Use of assert detected

(S101)


264-264: Use of assert detected

(S101)


265-265: Use of assert detected

(S101)


273-273: Use of assert detected

(S101)


274-274: Use of assert detected

(S101)

src/nat/agent/tool_calling_agent/agent.py

166-166: Redundant exception object included in logging.exception call

(TRY401)


192-192: Consider moving this statement to an else block

(TRY300)


194-194: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

🔇 Additional comments (3)
src/nat/agent/tool_calling_agent/agent.py (1)

26-26: LGTM: Needed import for graph construction.

tests/nat/agent/test_tool_calling.py (2)

254-257: LGTM: fixture provides compiled graph with return_direct.

The fixture is clear and aligns with the new API.


144-156: No action needed: pytest-asyncio is already set to asyncio_mode = "auto" in pyproject.toml, so async tests run automatically.

@thepatrickchin thepatrickchin changed the title Add return_direct option to tool_calling_agent to allow response to bypass LLM interpretation Add return_direct option to tool_calling_agent for direct tool responses Sep 9, 2025
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

🧹 Nitpick comments (11)
docs/source/workflows/about/tool-calling-agent.md (5)

89-90: Clarify behavior of return_direct vs. default post-processing.

Current text implies a “reasoning step” exists while the intro says tool-calling agents don’t reason between steps. Please phrase this as bypassing the final LLM post-processing/formatting step. Suggested tweak below.

-* `return_direct`: Optional list of tool names that should return their output directly without additional agent processing. When a tool in this list is called, its response is returned immediately to the user, bypassing the agent's reasoning step.
+* `return_direct`: Optional list of tool names whose outputs are returned as-is without LLM post-processing. When a listed tool is invoked, its raw result is returned immediately to the caller, and the workflow terminates.

77-91: Fix markdownlint MD004: use dash list markers for options.

markdownlint expects dashes. Update the entire “Configurable Options” list for consistency (including the newly added bullet).

-* `tool_names`: A list of tools that the agent can call. The tools must be functions configured in the YAML file
-
-* `llm_name`: The LLM the agent should use. The LLM must be configured in the YAML file
-
-* `verbose`: Defaults to False (useful to prevent logging of sensitive data). If set to True, the agent will log input, output, and intermediate steps.
-
-* `handle_tool_errors`: Defaults to True. All tool errors will be caught and a `ToolMessage` with an error message will be returned, allowing the agent to retry.
-
-* `max_iterations`: Defaults to 15. The maximum number of tool calls the agent may perform.
-
-* `return_direct`: Optional list of tool names that should return their output directly without additional agent processing. When a tool in this list is called, its response is returned immediately to the user, bypassing the agent's reasoning step.
-
-* `description`:  Defaults to "Tool Calling Agent Workflow". When the agent is configured as a function, this config option allows us to control the tool description (for example, when used as a tool within another agent).
+- `tool_names`: A list of tools that the agent can call. The tools must be functions configured in the YAML file
+
+- `llm_name`: The LLM the agent should use. The LLM must be configured in the YAML file
+
+- `verbose`: Defaults to False (useful to prevent logging of sensitive data). If set to True, the agent will log input, output, and intermediate steps.
+
+- `handle_tool_errors`: Defaults to True. All tool errors will be caught and a `ToolMessage` with an error message will be returned, allowing the agent to retry.
+
+- `max_iterations`: Defaults to 15. The maximum number of tool calls the agent may perform.
+
+- `return_direct`: Optional list of tool names whose outputs are returned as-is without LLM post-processing. When a listed tool is invoked, its raw result is returned immediately to the caller, and the workflow terminates.
+
+- `description`:  Defaults to "Tool Calling Agent Workflow". When the agent is configured as a function, this config option allows us to control the tool description (for example, when used as a tool within another agent).

100-114: Align “Response Handling” with return_direct semantics.

Section says the agent “passes” the tool result; by default it formats via LLM unless the tool is in return_direct. Add a short note.

-4. **Response Handling** – The tool returns a structured response, which the agent passes to the user.
+4. **Response Handling** – By default, the tool’s output is formatted by the LLM and then returned. If the tool is listed in `return_direct`, the raw tool output is returned immediately without LLM post-processing.

71-75: Add a minimal YAML example showing return_direct usage.

Helps readers adopt the new option quickly.

     verbose: true
     handle_tool_errors: true
+    return_direct:
+      - calculator_multiply
+      - calculator_divide
     description: 'Useful for performing simple mathematical calculations.'

20-21: Grammar fix for LLM support sentence.

Reads contradictory today. Suggested fix:

-Not all LLMs support tool calling / function calling, and can be used with tool calling agents.
+Not all LLMs support tool/function calling, so only those that do can be used with tool‑calling agents.
tests/nat/agent/test_tool_calling.py (6)

124-131: Rename fixture function to follow repo fixture naming convention

Guidelines ask fixtures to set name=... and the function itself to use the fixture_ prefix. Rename the function for consistency.

-@pytest.fixture(name='mock_tool_agent_with_return_direct', scope="module")
-def mock_agent_with_return_direct(mock_config_tool_calling_agent, mock_tool, mock_llm):
+@pytest.fixture(name='mock_tool_agent_with_return_direct', scope="module")
+def fixture_mock_tool_agent_with_return_direct(mock_config_tool_calling_agent, mock_tool, mock_llm):
     tools = [mock_tool('Tool A'), mock_tool('Tool B')]
     agent = ToolCallAgentGraph(llm=mock_llm,
                                tools=tools,
                                detailed_logs=mock_config_tool_calling_agent.verbose,
                                return_direct=['Tool A'])
     return agent

144-156: Avoid order fragility when asserting node keys

Graph node ordering can change; compare sets for stability.

-    assert list(graph.nodes.keys()) == ['__start__', 'agent', 'tool']
+    assert set(graph.nodes.keys()) == {'__start__', 'agent', 'tool'}

217-221: Empty messages relies on exception path—consider explicit guard in agent

Test is fine, but the underlying implementation currently hits an IndexError and then returns TOOL via the broad except. Prefer an explicit empty-state check to avoid exception-driven control flow.

Outside this file (in src/nat/agent/tool_calling_agent/agent.py), consider:

 async def tool_conditional_edge(self, state: ToolCallAgentGraphState):
-    try:
-        logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
-        last_message = state.messages[-1]
+    try:
+        logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
+        if not state.messages:
+            logger.debug("%s No messages in state; continuing to agent", AGENT_LOG_PREFIX)
+            return AgentDecision.TOOL
+        last_message = state.messages[-1]

254-257: Rename fixture function to follow fixture_ prefix

Same naming note as above for consistency.

-@pytest.fixture(name="mock_tool_graph_with_return_direct", scope="module")
-async def mock_graph_with_return_direct(mock_tool_agent_with_return_direct):
+@pytest.fixture(name="mock_tool_graph_with_return_direct", scope="module")
+async def fixture_mock_tool_graph_with_return_direct(mock_tool_agent_with_return_direct):
     return await mock_tool_agent_with_return_direct.build_graph()

268-275: Strengthen assertion for direct-return content

Asserting the content ensures the tool output itself is returned unchanged.

     last_message = response.messages[-1]
     assert isinstance(last_message, ToolMessage)
     assert last_message.name == 'Tool A'
+    assert last_message.content == 'mock query'

101-115: Ignore Ruff S101 for bare assert in test files
Add a rule to suppress S101 in your tests—choose one:

  • In pyproject.toml:
    [tool.ruff.lint.per-file-ignores]
    "tests/**" = ["S101"]
  • At the top of tests/nat/agent/test_tool_calling.py:
    # ruff: noqa: S101
📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 16c3762 and 1d91458.

📒 Files selected for processing (2)
  • docs/source/workflows/about/tool-calling-agent.md (1 hunks)
  • tests/nat/agent/test_tool_calling.py (4 hunks)
🧰 Additional context used
📓 Path-based instructions (8)
tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Unit tests must live under tests/ and use configured markers (e2e, integration, etc.)

Files:

  • tests/nat/agent/test_tool_calling.py

⚙️ CodeRabbit configuration file

tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using the test_ prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using the fixture_ prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass

Files:

  • tests/nat/agent/test_tool_calling.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

Files:

  • tests/nat/agent/test_tool_calling.py
**/tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/tests/**/*.py: Test functions must use the test_ prefix and snake_case
Extract repeated test code into pytest fixtures; fixtures should set name=... in @pytest.fixture and functions named with fixture_ prefix
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Use pytest with pytest-asyncio for async code; mock external services with pytest_httpserver or unittest.mock

Files:

  • tests/nat/agent/test_tool_calling.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • tests/nat/agent/test_tool_calling.py
  • docs/source/workflows/about/tool-calling-agent.md
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • tests/nat/agent/test_tool_calling.py
  • docs/source/workflows/about/tool-calling-agent.md
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • tests/nat/agent/test_tool_calling.py
  • docs/source/workflows/about/tool-calling-agent.md
docs/source/**/*.md

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

docs/source/**/*.md: Use the official naming: first use “NVIDIA NeMo Agent toolkit”; subsequent uses “NeMo Agent toolkit”; never use deprecated names in documentation
Documentation sources must be Markdown under docs/source; keep docs in sync and fix Sphinx errors/broken links
Documentation must be clear, comprehensive, free of TODO/FIXME/placeholders/offensive/outdated terms; fix spelling; adhere to Vale vocab allow/reject lists

Files:

  • docs/source/workflows/about/tool-calling-agent.md
docs/source/**/*

⚙️ CodeRabbit configuration file

This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the docs/source/_static directory.

Files:

  • docs/source/workflows/about/tool-calling-agent.md
🧬 Code graph analysis (1)
tests/nat/agent/test_tool_calling.py (3)
tests/conftest.py (2)
  • mock_llm (274-332)
  • mock_tool (336-358)
src/nat/agent/tool_calling_agent/agent.py (4)
  • ToolCallAgentGraph (47-204)
  • build_graph (197-204)
  • ToolCallAgentGraphState (42-44)
  • tool_conditional_edge (152-168)
src/nat/agent/base.py (1)
  • AgentDecision (62-64)
🪛 Ruff (0.12.2)
tests/nat/agent/test_tool_calling.py

108-108: Use of assert detected

(S101)


109-109: Use of assert detected

(S101)


110-110: Use of assert detected

(S101)


111-111: Use of assert detected

(S101)


112-112: Use of assert detected

(S101)


113-113: Use of assert detected

(S101)


114-114: Use of assert detected

(S101)


146-146: Use of assert detected

(S101)


147-147: Use of assert detected

(S101)


148-148: Use of assert detected

(S101)


149-149: Use of assert detected

(S101)


153-153: Use of assert detected

(S101)


154-154: Use of assert detected

(S101)


155-155: Use of assert detected

(S101)


193-193: Use of assert detected

(S101)


200-200: Use of assert detected

(S101)


207-207: Use of assert detected

(S101)


214-214: Use of assert detected

(S101)


220-220: Use of assert detected

(S101)


264-264: Use of assert detected

(S101)


265-265: Use of assert detected

(S101)


273-273: Use of assert detected

(S101)


274-274: Use of assert detected

(S101)

🪛 markdownlint-cli2 (0.17.2)
docs/source/workflows/about/tool-calling-agent.md

89-89: Unordered list style
Expected: dash; Actual: asterisk

(MD004, ul-style)

🔇 Additional comments (5)
tests/nat/agent/test_tool_calling.py (5)

101-115: Init test for return_direct looks correct

Covers construction and properties; no issues spotted.


189-194: Behavior without return_direct is validated

Confirms default path returns to agent; good coverage.


196-201: Direct-return match path is validated

END on matching tool name is correctly asserted.


203-208: Non-matching tool falls back to agent

Covers the negative case; looks good.


210-215: No name attribute case handled

Validates graceful fallback; OK.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (4)
src/nat/agent/tool_calling_agent/agent.py (4)

19-27: Import ToolMessage to enable safe type-checks in tool edge.

tool_conditional_edge should isinstance(..., ToolMessage). Add the import.

-from langchain_core.messages import SystemMessage
+from langchain_core.messages import SystemMessage
+from langchain_core.messages import ToolMessage

53-63: Type-hint callbacks as optional and normalize return_direct once.

  • Pyright will flag callbacks without | None.
  • Normalize return_direct to a lowercase frozenset for O(1) membership and case-insensitive matches.
-        callbacks: list[AsyncCallbackHandler] = None,
+        callbacks: list[AsyncCallbackHandler] | None = None,
@@
-        self.return_direct = return_direct
+        self.return_direct = frozenset(n.lower() for n in return_direct) if return_direct else frozenset()

Also applies to: 91-95


115-130: Fix logging per TRY401 and add return type.

Remove redundant exception object from logger.exception and annotate the return type.

-    async def conditional_edge(self, state: ToolCallAgentGraphState):
+    async def conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
@@
-        except Exception as ex:
-            logger.exception("%s Failed to determine whether agent is calling a tool: %s", AGENT_LOG_PREFIX, ex)
+        except Exception:
+            logger.exception("%s Failed to determine whether agent is calling a tool", AGENT_LOG_PREFIX)
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
             return AgentDecision.END

208-217: Docstring return type is wrong; function returns str | None, not ChatPromptTemplate.

Correct to avoid misleading users and mismatched docs.

 def create_tool_calling_agent_prompt(config: "ToolCallAgentWorkflowConfig") -> str | None:
     """
     Create a Tool Calling Agent prompt from the config.
 
     Args:
         config (ToolCallAgentWorkflowConfig): The config to use for the prompt.
 
     Returns:
-        ChatPromptTemplate: The Tool Calling Agent prompt.
+        str | None: The Tool Calling Agent system prompt string, or None if no prompt is configured.
     """
♻️ Duplicate comments (2)
src/nat/agent/tool_calling_agent/agent.py (2)

171-197: Optional: TRY300 – move return out of the try block; keep error+re-raise policy.

Ruff prefers returning in an else. Keep logger.error(...); raise per repo guidelines (do not switch to exception).

-    async def _build_graph(self, state_schema) -> CompiledGraph:
-        try:
+    async def _build_graph(self, state_schema) -> CompiledGraph:
+        try:
             logger.debug("%s Building and compiling the Tool Calling Agent Graph", AGENT_LOG_PREFIX)
@@
             self.graph = graph.compile()
-
-            return self.graph
-        except Exception as ex:
+        except Exception as ex:
             logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
             raise
+        else:
+            return self.graph

153-169: Harden tool edge: empty-state guard, type-check ToolMessage, case-insensitive match, and fix TRY401.

Prevents IndexError on empty state, avoids false positives on non-tool messages, and cleans logging per repo rules. Also add return type.

-    async def tool_conditional_edge(self, state: ToolCallAgentGraphState):
+    async def tool_conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
         try:
             logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
-            last_message = state.messages[-1]
-
-            if (self.return_direct and hasattr(last_message, 'name') and (last_message.name in self.return_direct)):
+            if not state.messages:
+                logger.debug("%s No messages in state; routing to agent", AGENT_LOG_PREFIX)
+                return AgentDecision.TOOL
+            last_message = state.messages[-1]
+            if (
+                self.return_direct
+                and isinstance(last_message, ToolMessage)
+                and last_message.name
+                and last_message.name.lower() in self.return_direct
+            ):
                 # Return directly if this tool is in the return_direct list
                 logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
                 return AgentDecision.END
             else:
                 # Continue to agent for processing
                 logger.debug("%s Tool response will be processed by agent", AGENT_LOG_PREFIX)
                 return AgentDecision.TOOL
-        except Exception as ex:
-            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+        except Exception:
+            logger.exception("%s Failed to determine tool conditional edge", AGENT_LOG_PREFIX)
             logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
             return AgentDecision.TOOL
🧹 Nitpick comments (2)
src/nat/agent/tool_calling_agent/agent.py (2)

198-205: Optional: TRY300 – use else for the return.

Minor style consistency with _build_graph.

-    async def build_graph(self) -> CompiledGraph:
-        try:
-            await self._build_graph(state_schema=ToolCallAgentGraphState)
-            logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
-            return self.graph
-        except Exception as ex:
-            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
-            raise
+    async def build_graph(self) -> CompiledGraph:
+        try:
+            await self._build_graph(state_schema=ToolCallAgentGraphState)
+            logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
+        except Exception as ex:
+            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
+        else:
+            return self.graph

48-52: Document the new return_direct behavior in the class docstring.

Add a brief note describing how return_direct short-circuits to END for specified tools.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 1d91458 and 1359e4d.

📒 Files selected for processing (1)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
🧰 Additional context used
📓 Path-based instructions (6)
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
🧬 Code graph analysis (1)
src/nat/agent/tool_calling_agent/agent.py (4)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/rewoo_agent/agent.py (3)
  • _build_graph (344-367)
  • conditional_edge (327-342)
  • build_graph (369-376)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
src/nat/agent/react_agent/agent.py (4)
  • agent_node (137-240)
  • tool_node (262-337)
  • conditional_edge (242-260)
  • build_graph (339-346)
🪛 Ruff (0.12.2)
src/nat/agent/tool_calling_agent/agent.py

167-167: Redundant exception object included in logging.exception call

(TRY401)


193-193: Consider moving this statement to an else block

(TRY300)


195-195: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

♻️ Duplicate comments (1)
src/nat/agent/tool_calling_agent/agent.py (1)

154-176: Fix redundant exception logging in tool_conditional_edge.

The method passes the exception object to logger.exception(), which is redundant since exception() automatically includes the traceback. This triggers TRY401.

Apply this diff to fix the logging:

         except Exception as ex:
-            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+            logger.exception("%s Failed to determine tool conditional edge", AGENT_LOG_PREFIX)
             logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
             return AgentDecision.TOOL
🧹 Nitpick comments (2)
src/nat/agent/tool_calling_agent/agent.py (2)

177-203: Consider restructuring exception handling in _build_graph.

The exception handling pattern could be improved:

  1. Line 199 (return statement) should be in an else block for cleaner control flow (TRY300)
  2. Line 201 uses logger.error() when re-raising, which is correct per your guidelines, but static analysis suggests logger.exception() (TRY400 - this is a false positive given your exception handling standards)

Apply this diff to improve the control flow:

     async def _build_graph(self, state_schema) -> CompiledGraph:
         try:
             logger.debug("%s Building and compiling the Tool Calling Agent Graph", AGENT_LOG_PREFIX)
 
             graph = StateGraph(state_schema)
             graph.add_node("agent", self.agent_node)
             graph.add_node("tool", self.tool_node)
 
             if self.return_direct:
                 # go to end of graph if tool is set to return directly
                 tool_conditional_edge_possible_outputs = {AgentDecision.END: "__end__", AgentDecision.TOOL: "agent"}
                 graph.add_conditional_edges("tool", self.tool_conditional_edge, tool_conditional_edge_possible_outputs)
             else:
                 # otherwise return to agent after tool call
                 graph.add_edge("tool", "agent")
 
             conditional_edge_possible_outputs = {AgentDecision.TOOL: "tool", AgentDecision.END: "__end__"}
             graph.add_conditional_edges("agent", self.conditional_edge, conditional_edge_possible_outputs)
 
             graph.set_entry_point("agent")
             self.graph = graph.compile()
-
-            return self.graph
         except Exception as ex:
             logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
             raise
+        else:
+            return self.graph

154-176: Add case-insensitive tool name matching for robustness.

The tool name comparison is case-sensitive, which could cause issues if tool names differ in casing between definition and runtime.

Consider normalizing both sides of the comparison:

             last_message = state.messages[-1]
             # Return directly if this tool is in the return_direct set
-            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
-                    and last_message.name in self.return_direct):
+            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name):
+                # Normalize for case-insensitive comparison
+                normalized_return_direct = {name.lower() for name in self.return_direct}
+                if last_message.name.lower() in normalized_return_direct:
-                # Return directly if this tool is in the return_direct list
-                logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
-                return AgentDecision.END
+                    # Return directly if this tool is in the return_direct list
+                    logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
+                    return AgentDecision.END
📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 1359e4d and 129fc03.

📒 Files selected for processing (1)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
🧰 Additional context used
📓 Path-based instructions (6)
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
🧬 Code graph analysis (1)
src/nat/agent/tool_calling_agent/agent.py (4)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/rewoo_agent/agent.py (3)
  • _build_graph (344-367)
  • conditional_edge (327-342)
  • build_graph (369-376)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
🪛 Ruff (0.12.2)
src/nat/agent/tool_calling_agent/agent.py

173-173: Redundant exception object included in logging.exception call

(TRY401)


199-199: Consider moving this statement to an else block

(TRY300)


201-201: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

Copy link
Member

@willkill07 willkill07 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good! Just some feedback on being more specific with types or generalizing by being able to pass a Pattern instead of a str for return_direct components

@willkill07
Copy link
Member

/ok to test 129fc03

Introduced `return_direct` field in ToolCallAgentWorkflowConfig to specify tools that return responses directly without LLM processing.

Signed-off-by: Patrick Chin <[email protected]>
Introduced `return_direct` field in ToolCallAgentWorkflowConfig to specify tools that return responses directly without LLM processing.

Signed-off-by: Patrick Chin <[email protected]>
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

♻️ Duplicate comments (1)
src/nat/agent/tool_calling_agent/agent.py (1)

173-175: Don’t pass ex to logger.exception (TRY401).

logger.exception already includes the traceback; remove the redundant argument.

-        except Exception as ex:
-            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+        except Exception:
+            logger.exception("%s Failed to determine tool conditional edge", AGENT_LOG_PREFIX)
             logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
             return AgentDecision.TOOL
🧹 Nitpick comments (6)
docs/source/workflows/about/tool-calling-agent.md (1)

89-90: Fix list bullet style to satisfy markdownlint MD004.

Use dashes for unordered lists per repo lint config.

-* `return_direct`: Optional list of tool names that should return their output directly without additional agent processing. When a tool in this list is called, its response is returned immediately to the user, bypassing the agent's reasoning step.
+- `return_direct`: Optional list of tool names that should return their output directly without additional agent processing. When a tool in this list is called, its response is returned immediately to the user, bypassing the agent's reasoning step.
src/nat/agent/tool_calling_agent/agent.py (4)

63-64: Broaden type for return_direct to accept any sequence (not just list).

Improves API ergonomics without breaking behavior.

-        return_direct: list[str] | None = None,
+        return_direct: typing.Sequence[str] | None = None,

94-94: Normalize and validate return_direct tool names; keep original attribute for introspection.

  • Make membership O(1) and case-insensitive.
  • Warn on unknown tool names to catch config typos.
-        self.return_direct = return_direct
+        self.return_direct = list(return_direct) if return_direct else None
+        # Build normalized lookup set and validate entries
+        self._return_direct_names: set[str] | None = None
+        if self.return_direct:
+            available = {t.name.lower() for t in tools}
+            requested = {name.lower() for name in self.return_direct}
+            invalid = requested - available
+            if invalid:
+                logger.warning("%s Tools specified in return_direct not found: %s", AGENT_LOG_PREFIX, sorted(invalid))
+            self._return_direct_names = requested & available

154-176: Add a concise docstring for tool_conditional_edge.

Helps readers and satisfies docstring expectations for public methods.

-    async def tool_conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
+    async def tool_conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
+        """Decide whether to end or route back to the agent after a tool call.
+
+        Returns:
+            AgentDecision.END if the last ToolMessage came from a tool listed in
+            `return_direct`; otherwise AgentDecision.TOOL to continue to the agent.
+        """

163-165: Compare tool names case-insensitively and use the precomputed lookup set.

Prevents config/user casing mismatches and speeds membership checks.

-            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
-                    and last_message.name in self.return_direct):
+            if (
+                self._return_direct_names is not None
+                and isinstance(last_message, ToolMessage)
+                and last_message.name
+                and last_message.name.lower() in self._return_direct_names
+            ):
tests/nat/agent/test_tool_calling.py (1)

144-156: Be cautious asserting internal builder edges; prefer behavior-focused checks.

LangGraph internals can shift; consider asserting execution outcomes or branch keys only.

-    assert graph.builder.edges == {('__start__', 'agent')}
+    # Prefer checking exposed branch keys and reachable nodes rather than raw edges
+    assert set(graph.builder.branches.get('agent').keys()) == {'conditional_edge'}
📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 129fc03 and fd254dc.

📒 Files selected for processing (4)
  • docs/source/workflows/about/tool-calling-agent.md (1 hunks)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
  • src/nat/agent/tool_calling_agent/register.py (2 hunks)
  • tests/nat/agent/test_tool_calling.py (4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/nat/agent/tool_calling_agent/register.py
🧰 Additional context used
📓 Path-based instructions (11)
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
  • docs/source/workflows/about/tool-calling-agent.md
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
  • docs/source/workflows/about/tool-calling-agent.md
**/*.{py,yaml,yml}

📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)

**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
  • docs/source/workflows/about/tool-calling-agent.md
tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Unit tests must live under tests/ and use configured markers (e2e, integration, etc.)

Files:

  • tests/nat/agent/test_tool_calling.py

⚙️ CodeRabbit configuration file

tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using the test_ prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using the fixture_ prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass

Files:

  • tests/nat/agent/test_tool_calling.py
**/tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/tests/**/*.py: Test functions must use the test_ prefix and snake_case
Extract repeated test code into pytest fixtures; fixtures should set name=... in @pytest.fixture and functions named with fixture_ prefix
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Use pytest with pytest-asyncio for async code; mock external services with pytest_httpserver or unittest.mock

Files:

  • tests/nat/agent/test_tool_calling.py
docs/source/**/*.md

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

docs/source/**/*.md: Use the official naming: first use “NVIDIA NeMo Agent toolkit”; subsequent uses “NeMo Agent toolkit”; never use deprecated names in documentation
Documentation sources must be Markdown under docs/source; keep docs in sync and fix Sphinx errors/broken links
Documentation must be clear, comprehensive, free of TODO/FIXME/placeholders/offensive/outdated terms; fix spelling; adhere to Vale vocab allow/reject lists

Files:

  • docs/source/workflows/about/tool-calling-agent.md
docs/source/**/*

⚙️ CodeRabbit configuration file

This directory contains the source code for the documentation. All documentation should be written in Markdown format. Any image files should be placed in the docs/source/_static directory.

Files:

  • docs/source/workflows/about/tool-calling-agent.md
🧬 Code graph analysis (2)
src/nat/agent/tool_calling_agent/agent.py (3)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
tests/nat/agent/test_tool_calling.py (3)
tests/conftest.py (2)
  • mock_llm (274-332)
  • mock_tool (336-358)
src/nat/agent/tool_calling_agent/agent.py (4)
  • ToolCallAgentGraph (49-211)
  • build_graph (204-211)
  • ToolCallAgentGraphState (44-46)
  • tool_conditional_edge (154-175)
src/nat/agent/base.py (1)
  • AgentDecision (62-64)
🪛 Ruff (0.12.2)
src/nat/agent/tool_calling_agent/agent.py

173-173: Redundant exception object included in logging.exception call

(TRY401)


199-199: Consider moving this statement to an else block

(TRY300)


201-201: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

🪛 markdownlint-cli2 (0.17.2)
docs/source/workflows/about/tool-calling-agent.md

89-89: Unordered list style
Expected: dash; Actual: asterisk

(MD004, ul-style)

🔇 Additional comments (7)
src/nat/agent/tool_calling_agent/agent.py (3)

22-29: Imports look correct for current LangGraph versions.

Using ToolMessage and CompiledStateGraph from the updated paths is appropriate.


177-199: Graph wiring for return_direct is sound.

Conditional edges on the tool node map END→__end__ and TOOL→agent as intended; entry point remains agent. Good separation from the non-return_direct path.


204-208: LGTM.

Return type annotation added and compiled graph build flow is clear.

tests/nat/agent/test_tool_calling.py (4)

101-115: Constructor test for return_direct is solid.

Asserts the attribute is plumbed and tools are registered. Good coverage.


124-131: Fixture for return_direct agent is clear and reusable.

Good use of module-scoped fixture to speed tests.


189-221: Decision-path tests cover key cases (match, no-match, no-name, empty state).

Nice coverage of edge conditions for tool_conditional_edge.


268-275: End-to-end direct-return assertion is correct.

Verifies the final message is a ToolMessage from the direct-return tool.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
src/nat/agent/tool_calling_agent/agent.py (1)

224-233: Docstring return type is incorrect (says ChatPromptTemplate but returns str | None).

Update to avoid confusion and doc-test drift.

 def create_tool_calling_agent_prompt(config: "ToolCallAgentWorkflowConfig") -> str | None:
-    """
-    Create a Tool Calling Agent prompt from the config.
-
-    Args:
-        config (ToolCallAgentWorkflowConfig): The config to use for the prompt.
-
-    Returns:
-        ChatPromptTemplate: The Tool Calling Agent prompt.
-    """
+    """
+    Create a Tool Calling Agent prompt string from the workflow config.
+
+    Args:
+        config (ToolCallAgentWorkflowConfig): The workflow config.
+
+    Returns:
+        str | None: The system prompt to prepend, or None if no prompt is configured.
+    """
♻️ Duplicate comments (4)
src/nat/agent/tool_calling_agent/agent.py (4)

127-129: Fix TRY401: don’t pass exception object to logger.exception.

Keep traceback via logger.exception without formatting ex.

-        except Exception as ex:
-            logger.exception("%s Failed to determine whether agent is calling a tool: %s", AGENT_LOG_PREFIX, ex)
+        except Exception:
+            logger.exception("%s Failed to determine whether agent is calling a tool", AGENT_LOG_PREFIX)
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)

63-64: return_direct bug: membership check will never match; normalize to tool-name set.

You accept list[BaseTool] but later compare a str (last_message.name) against that list, so direct-return never triggers. Normalize to a lowercase set of names and validate against available tools.

@@
-        return_direct: list[BaseTool] | None = None,
+        return_direct: list[str] | list[BaseTool] | None = None,
@@
-        self.agent = prompt_runnable | self.bound_llm
-        self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
-        self.return_direct = return_direct
+        self.agent = prompt_runnable | self.bound_llm
+        self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
+        # Normalize return_direct to a lowercase set of tool names and warn on unknowns
+        self._return_direct_names: set[str] = set()
+        if return_direct:
+            for item in return_direct:
+                name = item.name if isinstance(item, BaseTool) else str(item)
+                if name:
+                    self._return_direct_names.add(name.lower())
+            available = {t.name.lower() for t in tools}
+            invalid = self._return_direct_names - available
+            if invalid:
+                logger.warning("%s Tools specified in return_direct not found: %s",
+                               AGENT_LOG_PREFIX, sorted(invalid))
         logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX)

Also applies to: 92-96


171-181: Wire normalized names into routing and guard case-insensitively.

Use the normalized set and lower() comparison; also gate graph behavior on that set.

@@
-            # Return directly if this tool is in the return_direct set
-            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
-                    and last_message.name in self.return_direct):
+            # Return directly if this tool is in the normalized return_direct name set
+            if (self._return_direct_names
+                and isinstance(last_message, ToolMessage)
+                and last_message.name
+                and last_message.name.lower() in self._return_direct_names):
                 # Return directly if this tool is in the return_direct list
                 logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
                 return AgentDecision.END
             else:
                 # Continue to agent for processing
                 logger.debug("%s Tool response will be processed by agent", AGENT_LOG_PREFIX)
                 return AgentDecision.TOOL
@@
-            if self.return_direct:
+            if self._return_direct_names:
                 # go to end of graph if tool is set to return directly
                 tool_conditional_edge_possible_outputs = {AgentDecision.END: "__end__", AgentDecision.TOOL: "agent"}
                 graph.add_conditional_edges("tool", self.tool_conditional_edge, tool_conditional_edge_possible_outputs)
             else:
                 # otherwise return to agent after tool call
                 graph.add_edge("tool", "agent")

Also applies to: 195-201


183-185: Fix TRY401: don’t pass exception object to logger.exception.

Same issue in tool_conditional_edge.

-        except Exception as ex:
-            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+        except Exception:
+            logger.exception("%s Failed to determine tool conditional edge", AGENT_LOG_PREFIX)
             logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
🧹 Nitpick comments (2)
src/nat/agent/tool_calling_agent/agent.py (2)

136-146: Avoid shadowing “tools”; prefer clearer name in logs.

Minor readability improvement.

-            tools = [tool.get("name") for tool in tool_calls]
+            tool_names = [tool.get("name") for tool in tool_calls]
@@
-                    self._log_tool_response(str(tools), str(tool_input), response.content)
+                    self._log_tool_response(str(tool_names), str(tool_input), response.content)

59-59: Optional: refine callbacks type hint.

Prefer Optional and/or collections.abc.Sequence.

-        callbacks: list[AsyncCallbackHandler] = None,
+        callbacks: list[AsyncCallbackHandler] | None = None,
📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between ac68663 and 2664357.

📒 Files selected for processing (2)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
  • src/nat/agent/tool_calling_agent/register.py (2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/nat/agent/tool_calling_agent/register.py
🧰 Additional context used
📓 Path-based instructions (7)
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,yaml,yml}

📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)

**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
🧬 Code graph analysis (1)
src/nat/agent/tool_calling_agent/agent.py (3)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
🪛 Ruff (0.12.2)
src/nat/agent/tool_calling_agent/agent.py

183-183: Redundant exception object included in logging.exception call

(TRY401)


209-209: Consider moving this statement to an else block

(TRY300)


211-211: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

🔇 Additional comments (6)
src/nat/agent/tool_calling_agent/agent.py (6)

27-29: Correct import updates for LangGraph.

Using StateGraph and CompiledStateGraph aligns with the recent LangGraph bump.


22-22: Appropriate ToolMessage import.

Importing ToolMessage from langchain_core is correct for ToolNode outputs.


154-164: Docstring addition for tool_conditional_edge looks good.

Clear, concise, and matches behavior.


195-205: Graph wiring LGTM.

Conditional edge on tool with END/agent mapping matches the feature intent; fallback edge preserved.


214-218: Return type on build_graph LGTM.

Matches CompiledStateGraph; consistent with imports and base contract.


167-169: Good guard for empty state.

Safe early route to agent when no messages are present.

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Signed-off-by: Patrick Chin <[email protected]>
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

♻️ Duplicate comments (1)
src/nat/agent/tool_calling_agent/agent.py (1)

63-64: Fix broken return_direct logic (string vs BaseTool mismatch) and make matching case-safe.

return_direct is typed as list[BaseTool], but later you compare ToolMessage.name (a str) against that list, so it never matches. Normalize to a set of tool names and compare case-insensitively; also update the graph condition to use the normalized set.

Apply:

@@
-        return_direct: list[BaseTool] | None = None,
+        return_direct: list[BaseTool] | None = None,
@@
-        self.return_direct = return_direct
+        # Normalize configured direct-return tools to a case-insensitive name set
+        self.return_direct_names: set[str] = {t.name.lower() for t in return_direct} if return_direct else set()
@@
-            # Return directly if this tool is in the return_direct set
-            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
-                    and last_message.name in self.return_direct):
+            # Return directly if this tool is configured to return directly
+            if (
+                self.return_direct_names
+                and isinstance(last_message, ToolMessage)
+                and last_message.name
+                and last_message.name.lower() in self.return_direct_names
+            ):
@@
-            if self.return_direct:
+            if self.return_direct_names:
                 # go to end of graph if tool is set to return directly
-                tool_conditional_edge_possible_outputs = {AgentDecision.END: "__end__", AgentDecision.TOOL: "agent"}
+                tool_conditional_edge_possible_outputs = {AgentDecision.END: "__end__", AgentDecision.TOOL: "agent"}
                 graph.add_conditional_edges("tool", self.tool_conditional_edge, tool_conditional_edge_possible_outputs)
             else:
                 # otherwise return to agent after tool call
                 graph.add_edge("tool", "agent")

Also applies to: 94-95, 171-177, 195-201

🧹 Nitpick comments (5)
src/nat/agent/tool_calling_agent/agent.py (5)

128-129: Remove redundant exception object in logger.exception calls (TRY401).

Passing ex duplicates info; logger.exception already includes traceback.

-            logger.exception("%s Failed to determine whether agent is calling a tool: %s", AGENT_LOG_PREFIX, ex)
+            logger.exception("%s Failed to determine whether agent is calling a tool", AGENT_LOG_PREFIX)
@@
-            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+            logger.exception("%s Failed to determine tool conditional edge", AGENT_LOG_PREFIX)

Also applies to: 183-185


59-60: Make callbacks explicitly optional.

Defaulting to None but annotating as list[...] can trip Pyright.

-        callbacks: list[AsyncCallbackHandler] = None,
+        callbacks: list[AsyncCallbackHandler] | None = None,

195-201: Optional: small readability tweak.

Consider renaming the mapping to tool_edge_outputs to mirror conditional_edge_possible_outputs, but not required.


224-233: Docstring return type is incorrect.

Function returns str | None, not ChatPromptTemplate. Update the docstring for accuracy.

-    Returns:
-        ChatPromptTemplate: The Tool Calling Agent prompt.
+    Returns:
+        str | None: The Tool Calling Agent system prompt string, if any.

154-186: Future enhancement: allow wildcard/regex for tool selection.

Supporting patterns (e.g., calculator_*) via fnmatch or compiled regex would make return_direct more ergonomic. Could also accept list[Pattern[str] | str].

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2664357 and daeb1ac.

📒 Files selected for processing (1)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
🧰 Additional context used
📓 Path-based instructions (7)
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.{py,yaml,yml}

📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)

**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
🧬 Code graph analysis (1)
src/nat/agent/tool_calling_agent/agent.py (3)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
🪛 Ruff (0.12.2)
src/nat/agent/tool_calling_agent/agent.py

183-183: Redundant exception object included in logging.exception call

(TRY401)


209-209: Consider moving this statement to an else block

(TRY300)


211-211: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

🔇 Additional comments (2)
src/nat/agent/tool_calling_agent/agent.py (2)

27-29: Correct import updates for LangGraph types.

Switch to CompiledStateGraph and new import paths looks good and matches the recent bump.


154-164: Docstring for tool_conditional_edge looks good.

Clear behavior and return semantics documented.

@willkill07
Copy link
Member

/ok to test daeb1ac

@thepatrickchin thepatrickchin requested a review from a team as a code owner September 16, 2025 03:18
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

♻️ Duplicate comments (2)
src/nat/agent/tool_calling_agent/agent.py (2)

97-115: Add return type annotations for public methods.

Required by repo rules; see prior comment.

Apply:

-    async def agent_node(self, state: ToolCallAgentGraphState):
+    async def agent_node(self, state: ToolCallAgentGraphState) -> ToolCallAgentGraphState:
@@
-    async def conditional_edge(self, state: ToolCallAgentGraphState):
+    async def conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
@@
-    async def tool_node(self, state: ToolCallAgentGraphState):
+    async def tool_node(self, state: ToolCallAgentGraphState) -> ToolCallAgentGraphState:

Also applies to: 116-131, 132-153


166-186: Make comparison case‑insensitive and fix TRY401 (don’t pass exception object to logger.exception).

Keeps behavior resilient and satisfies lint.

Apply:

-            logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
+            logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
             if not state.messages:
                 logger.debug("%s No messages in state; routing to agent", AGENT_LOG_PREFIX)
                 return AgentDecision.TOOL
 
             last_message = state.messages[-1]
             # Return directly if this tool is in the return_direct set
-            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
-                    and last_message.name in self.return_direct):
+            if (
+                self.return_direct
+                and isinstance(last_message, ToolMessage)
+                and isinstance(last_message.name, str)
+                and last_message.name.lower() in self.return_direct
+            ):
                 # Return directly if this tool is in the return_direct list
                 logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
                 return AgentDecision.END
             else:
                 # Continue to agent for processing
                 logger.debug("%s Tool response will be processed by agent", AGENT_LOG_PREFIX)
                 return AgentDecision.TOOL
         except Exception as ex:
-            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+            logger.exception("%s Failed to determine tool conditional edge", AGENT_LOG_PREFIX)
             logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
             return AgentDecision.TOOL
🧹 Nitpick comments (7)
tests/nat/agent/test_tool_calling.py (3)

124-131: Rename fixture function to follow repo fixture naming convention.

Per guidelines, fixture function names should use the fixture_ prefix (you’re already setting name=...).

Apply:

-@pytest.fixture(name='mock_tool_agent_with_return_direct', scope="module")
-def mock_agent_with_return_direct(mock_config_tool_calling_agent, mock_tool, mock_llm):
+@pytest.fixture(name='mock_tool_agent_with_return_direct', scope="module")
+def fixture_mock_tool_agent_with_return_direct(mock_config_tool_calling_agent, mock_tool, mock_llm):

144-156: Good graph assertions; avoid coupling to internal builder details if LangGraph changes.

If LangGraph internals change, these may become brittle. Prefer validating public behavior (node names, final outputs) when feasible.


254-257: Rename graph fixture function for consistency.

Match fixture function naming convention.

Apply:

-@pytest.fixture(name="mock_tool_graph_with_return_direct", scope="module")
-async def mock_graph_with_return_direct(mock_tool_agent_with_return_direct):
+@pytest.fixture(name="mock_tool_graph_with_return_direct", scope="module")
+async def fixture_mock_graph_with_return_direct(mock_tool_agent_with_return_direct):
src/nat/agent/tool_calling_agent/agent.py (4)

94-95: Normalize and validate return_direct for robustness (case-insensitive, O(1) lookup).

Prevents case bugs and surfaces config mistakes early.

Apply:

-        self.return_direct = [tool.name for tool in return_direct] if return_direct else []
+        # Normalize to lowercase set for stable, O(1) lookups
+        self.return_direct = {tool.name.lower() for tool in return_direct} if return_direct else set()
+        # Validate presence among provided tools (helps catch config typos)
+        if self.return_direct:
+            available_tools = {tool.name.lower() for tool in tools}
+            invalid = self.return_direct - available_tools
+            if invalid:
+                logger.warning(
+                    "%s Tools specified in return_direct not found among provided tools: %s",
+                    AGENT_LOG_PREFIX,
+                    sorted(invalid),
+                )

187-213: Keep logger.error + raise (per repo policy); silence Ruff TRY400 locally.

Add per-line noqa where re-raising to avoid linter-auto-fix to logger.exception.

Apply:

-        except Exception as ex:
-            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
+        except Exception as ex:
+            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)  # noqa: TRY400
             raise

214-221: Same note for build_graph: preserve logger.error + raise; add noqa.

Matches exception-logging policy while keeping CI green.

Apply:

-        except Exception as ex:
-            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
+        except Exception as ex:
+            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)  # noqa: TRY400
             raise

63-63: Optional: support glob/regex patterns for return_direct.

Aligns with maintainer suggestion; can map names once and check with fnmatch or compiled regexes.

Also applies to: 94-95, 166-176

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between daeb1ac and fe558c8.

📒 Files selected for processing (3)
  • src/nat/agent/tool_calling_agent/agent.py (4 hunks)
  • src/nat/agent/tool_calling_agent/register.py (2 hunks)
  • tests/nat/agent/test_tool_calling.py (4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/nat/agent/tool_calling_agent/register.py
🧰 Additional context used
📓 Path-based instructions (9)
src/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
src/nat/**/*

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Core functionality under src/nat should prioritize backward compatibility when changed

Files:

  • src/nat/agent/tool_calling_agent/agent.py

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bare raise and log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial

**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
**/*.{py,md}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
**/*.{py,yaml,yml}

📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)

**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

  • Ensure the code follows best practices and coding standards. - For Python code, follow
    PEP 20 and
    PEP 8 for style guidelines.
  • Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
    Example:
    def my_function(param1: int, param2: str) -> bool:
        pass
  • For Python exception handling, ensure proper stack trace preservation:
    • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
      and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
    • When catching and logging exceptions without re-raising: always use logger.exception()
      to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,

and should contain an Apache License 2.0 header comment at the top of each file.

  • Confirm that copyright years are up-to date whenever a file is changed.

Files:

  • src/nat/agent/tool_calling_agent/agent.py
  • tests/nat/agent/test_tool_calling.py
tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Unit tests must live under tests/ and use configured markers (e2e, integration, etc.)

Files:

  • tests/nat/agent/test_tool_calling.py

⚙️ CodeRabbit configuration file

tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using the test_ prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using the fixture_ prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass

Files:

  • tests/nat/agent/test_tool_calling.py
**/tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/tests/**/*.py: Test functions must use the test_ prefix and snake_case
Extract repeated test code into pytest fixtures; fixtures should set name=... in @pytest.fixture and functions named with fixture_ prefix
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Use pytest with pytest-asyncio for async code; mock external services with pytest_httpserver or unittest.mock

Files:

  • tests/nat/agent/test_tool_calling.py
🧬 Code graph analysis (2)
src/nat/agent/tool_calling_agent/agent.py (3)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
src/nat/agent/base.py (2)
  • AgentDecision (62-64)
  • _build_graph (259-260)
src/nat/agent/dual_node.py (4)
  • _build_graph (58-72)
  • agent_node (47-48)
  • tool_node (51-52)
  • conditional_edge (55-56)
tests/nat/agent/test_tool_calling.py (4)
tests/conftest.py (2)
  • mock_llm (274-332)
  • mock_tool (336-358)
src/nat/agent/tool_calling_agent/agent.py (4)
  • ToolCallAgentGraph (49-221)
  • build_graph (214-221)
  • ToolCallAgentGraphState (44-46)
  • tool_conditional_edge (154-185)
src/nat/agent/base.py (1)
  • AgentDecision (62-64)
packages/nvidia_nat_data_flywheel/src/nat/plugins/data_flywheel/observability/schema/sink/elasticsearch/dfw_es_record.py (1)
  • ToolMessage (88-93)
🪛 Ruff (0.12.2)
src/nat/agent/tool_calling_agent/agent.py

183-183: Redundant exception object included in logging.exception call

(TRY401)


209-209: Consider moving this statement to an else block

(TRY300)


211-211: Use logging.exception instead of logging.error

Replace with exception

(TRY400)

🔇 Additional comments (9)
tests/nat/agent/test_tool_calling.py (5)

101-115: Init wiring for return_direct looks correct.

Mapping BaseTool -> names and asserting ['Tool A'] is solid.


189-194: No-return_direct path verified (routes back to agent).

Behavior matches expectation.


196-208: return_direct match/no‑match branches covered.

Happy path and negative case both asserted.


210-221: Edge cases covered (non-ToolMessage, empty messages).

These protect against regressions in conditional routing.


268-275: ToolMessage(name=...) accepted by pinned langchain-core — no change required. API docs show ToolMessage signature includes name: Optional[str] = None (observed in v0.1.20 and v0.2.x), so the test assertion is compatible.

src/nat/agent/tool_calling_agent/agent.py (4)

22-22: Import ToolMessage is correct for type checks.


27-29: Updated imports match bumped LangGraph (CompiledStateGraph).


63-63: Public API: return_direct typed as list[BaseTool]—good.

Aligns with discussion and keeps strong typing at the edge.


154-164: Docstring reads well and clarifies routing semantics.

@willkill07
Copy link
Member

/ok to test 36a4859

@willkill07
Copy link
Member

/merge

@rapids-bot rapids-bot bot merged commit 95e12e6 into NVIDIA:develop Sep 16, 2025
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Allow tool calls to bypass LLM interpretation and return response directly

2 participants