-
Notifications
You must be signed in to change notification settings - Fork 487
Add return_direct option to tool_calling_agent for direct tool responses #775
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
WalkthroughAdds a configurable Changes
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
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing touches
🧪 Generate unit tests
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.
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. Comment |
There was a problem hiding this 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 = Falseand 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 testsAlternatively, 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 withoutnameto confirmtool_conditional_edgefalls back toAgentDecision.TOOL.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 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.mdsrc/nat/agent/tool_calling_agent/register.pytests/nat/agent/test_tool_calling.pysrc/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.mdsrc/nat/agent/tool_calling_agent/register.pytests/nat/agent/test_tool_calling.pysrc/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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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.mdsrc/nat/agent/tool_calling_agent/register.pytests/nat/agent/test_tool_calling.pysrc/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/_staticdirectory.
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.pysrc/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.pysrc/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.pysrc/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 bareraiseand 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.pytests/nat/agent/test_tool_calling.pysrc/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 thetest_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 thefixture_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 toasyncio_mode = "auto"in pyproject.toml, so async tests run automatically.
16c3762 to
1d91458
Compare
There was a problem hiding this 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 conventionGuidelines 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 keysGraph 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 agentTest 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_ prefixSame 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 contentAsserting 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 bareassertin 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
📒 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 thetest_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 thefixture_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 bareraiseand 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.pydocs/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.pydocs/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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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.pydocs/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/_staticdirectory.
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 correctCovers construction and properties; no issues spotted.
189-194: Behavior without return_direct is validatedConfirms default path returns to agent; good coverage.
196-201: Direct-return match path is validatedEND on matching tool name is correctly asserted.
203-208: Non-matching tool falls back to agentCovers the negative case; looks good.
210-215: No name attribute case handledValidates graceful fallback; OK.
There was a problem hiding this 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: ImportToolMessageto enable safe type-checks in tool edge.
tool_conditional_edgeshouldisinstance(..., 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-hintcallbacksas optional and normalizereturn_directonce.
- Pyright will flag
callbackswithout| None.- Normalize
return_directto a lowercasefrozensetfor 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.exceptionand 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 returnsstr | None, notChatPromptTemplate.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 – movereturnout of thetryblock; keep error+re-raise policy.Ruff prefers returning in an
else. Keeplogger.error(...); raiseper repo guidelines (do not switch toexception).- 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-checkToolMessage, case-insensitive match, and fix TRY401.Prevents
IndexErroron 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 – useelsefor 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 newreturn_directbehavior in the class docstring.Add a brief note describing how
return_directshort-circuits to END for specified tools.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 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 bareraiseand 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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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)
There was a problem hiding this 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 sinceexception()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:
- Line 199 (return statement) should be in an else block for cleaner control flow (TRY300)
- Line 201 uses
logger.error()when re-raising, which is correct per your guidelines, but static analysis suggestslogger.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
📒 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 bareraiseand 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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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)
willkill07
left a comment
There was a problem hiding this 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
|
/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]>
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]>
Signed-off-by: Patrick Chin <[email protected]>
Signed-off-by: Patrick Chin <[email protected]>
129fc03 to
fd254dc
Compare
There was a problem hiding this 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 passexto 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
📒 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 bareraiseand 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.pytests/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.pytests/nat/agent/test_tool_calling.pydocs/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.pytests/nat/agent/test_tool_calling.pydocs/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.pytests/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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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.pytests/nat/agent/test_tool_calling.pydocs/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 thetest_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 thefixture_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/_staticdirectory.
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.
Signed-off-by: Patrick Chin <[email protected]>
ac68663 to
2664357
Compare
There was a problem hiding this 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
📒 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 bareraiseand 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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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]>
There was a problem hiding this 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_directis typed aslist[BaseTool], but later you compareToolMessage.name(astr) 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
exduplicates info;logger.exceptionalready 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: Makecallbacksexplicitly optional.Defaulting to
Nonebut annotating aslist[...]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_outputsto mirrorconditional_edge_possible_outputs, but not required.
224-233: Docstring return type is incorrect.Function returns
str | None, notChatPromptTemplate. 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_*) viafnmatchor compiled regex would makereturn_directmore ergonomic. Could also acceptlist[Pattern[str] | str].
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 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 bareraiseand 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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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
CompiledStateGraphand 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.
|
/ok to test daeb1ac |
Signed-off-by: Patrick Chin <[email protected]>
There was a problem hiding this 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 settingname=...).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
fnmatchor compiled regexes.Also applies to: 94-95, 166-176
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 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 bareraiseand 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.pytests/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.pytests/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.pytests/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.pytests/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
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.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.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile 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.pytests/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 thetest_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 thefixture_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 includesname: 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.
|
/ok to test 36a4859 |
|
/merge |
Description
Closes #758
By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
Documentation
Tests