feat!: Unify async in Builder's get_* functions; make function group filters async
#834
Conversation
WalkthroughProject-wide migration converting many previously synchronous builder, workflow, function-group, tool retrieval, and related APIs to asynchronous (async/await). Examples, docs, front ends, agents, runtime, eval, experimental modules, profiler typings, tests, and test utilities were updated to call and mock the new async APIs. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Caller
participant WB as WorkflowBuilder
participant Reg as Registry
participant FG as FunctionGroup
participant Wrap as ToolWrapper
Caller->>WB: get_tools(refs, wrapper_type)
activate WB
WB->>WB: resolve refs concurrently (asyncio.gather)
par fetch functions/groups
WB->>Reg: fetch Function A
WB->>Reg: fetch FunctionGroup G
WB->>Reg: fetch Function B
and
Reg-->>WB: Function A
Reg-->>WB: FunctionGroup G
Reg-->>WB: Function B
end
WB->>FG: get_accessible_functions() (await)
FG-->>WB: dict[name, Function]
WB->>Wrap: build wrappers
Wrap-->>WB: list[Tool]
WB-->>Caller: list[Tool]
deactivate WB
sequenceDiagram
autonumber
actor FrontEnd
participant WB as WorkflowBuilder
participant SM as SessionManager
FrontEnd->>WB: build(entry_function?)
activate WB
WB->>WB: assemble workflow (await dependencies, functions/groups/tools)
WB-->>FrontEnd: Workflow (awaited)
deactivate WB
FrontEnd->>SM: create with Workflow
SM-->>FrontEnd: ready
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests
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 |
willkill07
added
improvement
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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)
examples/notebooks/1_getting_started.ipynb (1)
1-10: Add SPDX header to the notebookAll files must include the Apache-2.0 SPDX header. Add a first markdown cell containing the header comment.
Example cell to insert at index 0:
{ "cell_type": "markdown", "metadata": {}, "source": [ "<!--\\n", "SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\\n", "SPDX-License-Identifier: Apache-2.0\\n", "-->" ] }
🧹 Nitpick comments (18)
examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.py (2)
78-80: Cast plugin_name to str to avoid type mismatchesFunctionRef may not be a plain str; Semantic Kernel expects a string plugin_name.
Apply this diff:
- for tool_name, tool in zip(config.tool_names, tools): - kernel.add_plugin(plugin=tool, plugin_name=tool_name) + for tool_name, tool in zip(config.tool_names, tools): + kernel.add_plugin(plugin=tool, plugin_name=str(tool_name))
127-129: Don’t log GeneratorExit as an exceptionThis is a normal control-flow event; prefer debug/info to avoid noisy stack traces.
- except GeneratorExit: - logger.exception("Exited early!") + except GeneratorExit: + logger.info("Exited early")docs/source/workflows/function-groups.md (2)
20-22: Use official first mention: “NVIDIA NeMo Agent toolkit”First occurrence should include “NVIDIA”; later uses can say “NeMo Agent toolkit”.
-Function groups let you package multiple related functions together so they can share configuration, context, and resources within the NeMo Agent toolkit. +Function groups let you package multiple related functions together so they can share configuration, context, and resources within the NVIDIA NeMo Agent toolkit.
66-72: Fix import/name mismatch in exampleImport WorkflowBuilder (not Builder) to match later usage.
-from nat.builder.workflow_builder import Builder -from nat.builder.function import FunctionGroup +from nat.builder.workflow_builder import WorkflowBuilder +from nat.builder.function import FunctionGroupsrc/nat/agent/react_agent/register.py (1)
103-105: Clarify error messageReference missing tools, not the LLM name.
- if not tools: - raise ValueError(f"No tools specified for ReAct Agent '{config.llm_name}'") + if not tools: + raise ValueError(f"No tools resolved from tool_names: {config.tool_names}")examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (3)
82-89: Remove duplicate tool retrieval; bind onceYou manually build tools then fetch them again with get_tools. Keep a single awaited call and bind the result.
- # Get tools for alert triage - tool_names = config.tool_names - tools = [] - for tool_name in tool_names: - tool = builder.get_tool(tool_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN) - tools.append(tool) - llm_n_tools = llm.bind_tools(tools, parallel_tool_calls=True) + # Get tools for alert triage (async API) and bind once + tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + llm_n_tools = llm.bind_tools(tools, parallel_tool_calls=True) @@ - # Get tools specified in config - tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + # tools already resolved aboveAlso applies to: 104-109
54-55: Avoid mutable list default in Pydantic configUse default_factory to prevent shared mutable defaults.
- tool_names: list[str] = [] + tool_names: list[str] = Field(default_factory=list)
19-21: Import Field from pydantic (canonical path)Aligns with project convention and pydantic v2 style.
-from pydantic.fields import Field +from pydantic import Fieldsrc/nat/agent/rewoo_agent/register.py (1)
112-114: Clarify error messagePoint to tool_names for easier debugging.
- if not tools: - raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'") + if not tools: + raise ValueError(f"No tools resolved from tool_names: {config.tool_names}")examples/notebooks/1_getting_started.ipynb (1)
19-23: Update repo link to official NVIDIA orgPoint users to the canonical repository.
-3. NeMo-Agent-Toolkit installed from source following [these instructions](https://github.com/cdgamarose-nv/NeMo-Agent-Toolkit/tree/develop?tab=readme-ov-file#install-from-source) +3. NeMo-Agent-Toolkit installed from source following [these instructions](https://github.com/NVIDIA/NeMo-Agent-Toolkit/tree/develop?tab=readme-ov-file#install-from-source)examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.py (1)
142-145: Log exceptions with stack trace when not re-raisingUse logger.exception() since the error is handled locally.
- except Exception as e: - logger.error("Error in agno_personal_finance function: %s", str(e)) - return f"Sorry, I encountered an error while generating your financial plan: {str(e)}" + except Exception as e: + logger.exception("Error in agno_personal_finance function: %s", str(e)) + return f"Sorry, I encountered an error while generating your financial plan: {str(e)}"tests/nat/front_ends/mcp/test_mcp_custom_routes.py (1)
75-79: Style: moveAsyncMockimport to module scopeKeep imports top‑level with other mocks for consistency.
-from unittest.mock import Mock -from unittest.mock import patch +from unittest.mock import Mock +from unittest.mock import patch +from unittest.mock import AsyncMock ... - from unittest.mock import AsyncMocksrc/nat/runtime/loader.py (1)
115-118: Naming nit: variableworkflowholds aWorkflowBuilderMinor clarity tweak: rename the context var to
builder.- async with WorkflowBuilder.from_config(config=config) as workflow: - yield SessionManager(await workflow.build(), max_concurrency=max_concurrency) + async with WorkflowBuilder.from_config(config=config) as builder: + yield SessionManager(await builder.build(), max_concurrency=max_concurrency)packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
179-181: Align signature with abstract method and silence ARG002.Match Builder.get_tools type hints and prefix unused params to satisfy ruff.
Apply this diff:
- async def get_tools(self, tool_names: Sequence[str], wrapper_type): - """Mock implementation.""" - return [] + async def get_tools( + self, + _tool_names: Sequence[str | FunctionRef | FunctionGroupRef], + _wrapper_type: LLMFrameworkEnum | str, + ) -> list[typing.Any]: + """Mock implementation.""" + return []Also add the missing imports near the other imports:
from nat.builder.framework_enum import LLMFrameworkEnum from nat.data_models.component_ref import FunctionRef, FunctionGroupReftests/nat/mcp/test_mcp_client_impl.py (1)
108-126: Mark async tests with pytest-asyncio.These async tests should be marked to run under pytest-asyncio.
Apply this diff to each test:
+@pytest.mark.asyncio async def test_mcp_client_function_group_includes_respected():+@pytest.mark.asyncio async def test_mcp_client_function_group_applies_overrides():+@pytest.mark.asyncio async def test_mcp_client_function_group_no_include_exposes_all():And ensure
import pytestis present at the top of the file.Also applies to: 128-146, 158-161
src/nat/agent/tool_calling_agent/register.py (1)
63-70: Awaiting get_tools is correct; tweak exception for ruff TRY003.Shorten the error message to satisfy TRY003.
Apply this diff:
- if not tools: - raise ValueError(f"No tools specified for Tool Calling Agent '{config.llm_name}'") + if not tools: + raise ValueError("No tools specified for Tool Calling Agent.")src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py (1)
240-242: Build the workflow once to avoid duplicate build calls.Minor efficiency/readability win.
Apply this diff:
- await self.add_default_route(app, SessionManager(await builder.build())) - await self.add_evaluate_route(app, SessionManager(await builder.build())) + workflow = await builder.build() + await self.add_default_route(app, SessionManager(workflow)) + await self.add_evaluate_route(app, SessionManager(workflow)) @@ - entry_workflow = await builder.build(entry_function=ep.function_name) + entry_workflow = await builder.build(entry_function=ep.function_name)Also applies to: 247-249
src/nat/builder/function.py (1)
458-458: Extract repeated "Unknown ... functions" error messages to a constant.
Occurrences: src/nat/builder/function.py:458, 550, 609. Replace repeated f-strings with a single template or two small constants; e.g.:_UNKNOWN_FUNCTIONS_ERROR = "Unknown {kind} functions: {items}"
raise ValueError(_UNKNOWN_FUNCTIONS_ERROR.format(kind="excluded", items=sorted(missing)))
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (30)
docs/source/extend/function-groups.md(4 hunks)docs/source/workflows/function-groups.md(2 hunks)examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py(1 hunks)examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.py(1 hunks)examples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.py(1 hunks)examples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.py(1 hunks)examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.py(1 hunks)examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.py(1 hunks)examples/notebooks/1_getting_started.ipynb(1 hunks)examples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.py(1 hunks)examples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py(1 hunks)packages/nvidia_nat_test/src/nat/test/tool_test_runner.py(1 hunks)src/nat/agent/react_agent/register.py(1 hunks)src/nat/agent/rewoo_agent/register.py(1 hunks)src/nat/agent/tool_calling_agent/register.py(1 hunks)src/nat/builder/builder.py(1 hunks)src/nat/builder/function.py(12 hunks)src/nat/builder/workflow_builder.py(6 hunks)src/nat/control_flow/router_agent/register.py(1 hunks)src/nat/control_flow/sequential_executor.py(1 hunks)src/nat/eval/evaluate.py(1 hunks)src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py(1 hunks)src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py(3 hunks)src/nat/front_ends/simple_base/simple_front_end_plugin_base.py(1 hunks)src/nat/runtime/loader.py(1 hunks)tests/nat/builder/test_builder.py(11 hunks)tests/nat/builder/test_function_group.py(22 hunks)tests/nat/front_ends/mcp/test_mcp_custom_routes.py(1 hunks)tests/nat/front_ends/mcp/test_mcp_front_end_plugin.py(8 hunks)tests/nat/mcp/test_mcp_client_impl.py(3 hunks)
🧰 Additional context used
📓 Path-based instructions (14)
**/*.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:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.pysrc/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pyexamples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.pyexamples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.pyexamples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pyexamples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pytests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pytests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pytests/nat/builder/test_builder.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:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.pysrc/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pyexamples/notebooks/1_getting_started.ipynbexamples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.pyexamples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.pyexamples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pyexamples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pydocs/source/extend/function-groups.mdtests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pydocs/source/workflows/function-groups.mdtests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pytests/nat/builder/test_builder.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:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.pysrc/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pyexamples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.pyexamples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.pyexamples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pyexamples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pydocs/source/extend/function-groups.mdtests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pydocs/source/workflows/function-groups.mdtests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pytests/nat/builder/test_builder.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:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.pysrc/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pyexamples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.pyexamples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.pyexamples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pyexamples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pytests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pytests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pytests/nat/builder/test_builder.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:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.pysrc/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pyexamples/notebooks/1_getting_started.ipynbexamples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.pyexamples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.pyexamples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pyexamples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pydocs/source/extend/function-groups.mdtests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pydocs/source/workflows/function-groups.mdtests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pytests/nat/builder/test_builder.py
examples/**/*
⚙️ CodeRabbit configuration file
examples/**/*: - This directory contains example code and usage scenarios for the toolkit, at a minimum an example should
contain a README.md or file README.ipynb.
- If an example contains Python code, it should be placed in a subdirectory named
src/and should
contain apyproject.tomlfile. Optionally, it might also contain scripts in ascripts/directory.- If an example contains YAML files, they should be placed in a subdirectory named
configs/. - If an example contains sample data files, they should be placed in a subdirectory nameddata/, and should
be checked into git-lfs.
Files:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.pyexamples/notebooks/1_getting_started.ipynbexamples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.pyexamples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.pyexamples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.pyexamples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.pyexamples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py
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/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.py
src/nat/**/*
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Core functionality under src/nat should prioritize backward compatibility when changed
Files:
src/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.py
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/builder/builder.pysrc/nat/control_flow/router_agent/register.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pysrc/nat/agent/react_agent/register.pysrc/nat/agent/rewoo_agent/register.pysrc/nat/agent/tool_calling_agent/register.pysrc/nat/eval/evaluate.pysrc/nat/runtime/loader.pysrc/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.pysrc/nat/builder/workflow_builder.pysrc/nat/control_flow/sequential_executor.pysrc/nat/builder/function.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.py
packages/*/src/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
packages/*/src/**/*.py: All importable Python code in packages must live under packages//src/
All public APIs in packaged code require Python 3.11+ type hints; prefer typing/collections.abc; use typing.Annotated when useful
Provide Google-style docstrings for public APIs in packages; first line concise with a period; use backticks for code entities
Files:
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py
packages/**/*
⚙️ CodeRabbit configuration file
packages/**/*: - This directory contains optional plugin packages for the toolkit, each should contain apyproject.tomlfile. - Thepyproject.tomlfile should declare a dependency onnvidia-nator another package with a name starting
withnvidia-nat-. This dependency should be declared using~=<version>, and the version should be a two
digit version (ex:~=1.0).
- Not all packages contain Python code, if they do they should also contain their own set of tests, in a
tests/directory at the same level as thepyproject.tomlfile.
Files:
packages/nvidia_nat_test/src/nat/test/tool_test_runner.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/extend/function-groups.mddocs/source/workflows/function-groups.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/extend/function-groups.mddocs/source/workflows/function-groups.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/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pytests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_builder.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/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pytests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_builder.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/mcp/test_mcp_client_impl.pytests/nat/builder/test_function_group.pytests/nat/front_ends/mcp/test_mcp_custom_routes.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_builder.py
🧠 Learnings (1)
📚 Learning: 2025-09-09T20:32:39.016Z
Learnt from: CR
PR: NVIDIA/NeMo-Agent-Toolkit#0
File: .cursor/rules/nat-test-llm.mdc:0-0
Timestamp: 2025-09-09T20:32:39.016Z
Learning: Applies to **/*.py : When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.<FRAMEWORK>) and call the framework’s method (e.g., ainvoke, achat, call).
Applied to files:
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.pyexamples/notebooks/1_getting_started.ipynbsrc/nat/agent/react_agent/register.pysrc/nat/agent/tool_calling_agent/register.pyexamples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py
🧬 Code graph analysis (26)
examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/builder/builder.py (4)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)tests/nat/mcp/test_mcp_client_impl.py (1)
get_tools(73-75)src/nat/data_models/component_ref.py (2)
FunctionRef(94-102)FunctionGroupRef(105-113)
src/nat/control_flow/router_agent/register.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/front_ends/simple_base/simple_front_end_plugin_base.py (2)
src/nat/runtime/session.py (1)
workflow(84-85)src/nat/builder/workflow_builder.py (1)
build(217-332)
examples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.py (4)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tool(183-185)src/nat/builder/builder.py (1)
get_tool(116-117)src/nat/builder/workflow_builder.py (2)
get_tool(574-588)get_tool(1181-1187)tests/nat/mcp/test_mcp_client_impl.py (1)
get_tool(69-71)
examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.py (2)
src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
examples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.py (4)
examples/object_store/user_report/tests/test_user_report_tool.py (1)
builder(33-49)packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/agent/react_agent/register.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/agent/rewoo_agent/register.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/agent/tool_calling_agent/register.py (2)
src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/eval/evaluate.py (3)
src/nat/runtime/session.py (2)
workflow(84-85)SessionManager(45-173)src/nat/builder/workflow_builder.py (1)
build(217-332)tests/nat/eval/test_evaluate.py (1)
session_manager(172-201)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (3)
src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)tests/nat/mcp/test_mcp_client_impl.py (1)
get_tools(73-75)
examples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/runtime/loader.py (2)
src/nat/runtime/session.py (2)
SessionManager(45-173)workflow(84-85)src/nat/builder/workflow_builder.py (1)
build(217-332)
src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py (2)
src/nat/runtime/session.py (1)
SessionManager(45-173)src/nat/builder/workflow_builder.py (1)
build(217-332)
tests/nat/mcp/test_mcp_client_impl.py (1)
src/nat/builder/function.py (1)
get_accessible_functions(485-521)
tests/nat/builder/test_function_group.py (1)
src/nat/builder/function.py (8)
get_accessible_functions(485-521)get_all_functions(628-664)add_function(380-428)get_included_functions(580-626)get_excluded_functions(523-578)set_filter_fn(666-675)set_per_function_filter_fn(677-695)_fn_should_be_included(444-447)
tests/nat/front_ends/mcp/test_mcp_custom_routes.py (2)
src/nat/builder/workflow_builder.py (1)
build(217-332)src/nat/builder/function.py (1)
get_accessible_functions(485-521)
tests/nat/front_ends/mcp/test_mcp_front_end_plugin.py (3)
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (2)
_get_all_functions(89-110)MCPFrontEndPluginWorker(210-251)src/nat/front_ends/mcp/mcp_front_end_plugin.py (2)
MCPFrontEndPlugin(28-113)_get_worker_instance(46-57)src/nat/builder/function.py (1)
get_accessible_functions(485-521)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/builder/workflow_builder.py (2)
src/nat/builder/function.py (2)
get_included_functions(580-626)get_accessible_functions(485-521)src/nat/builder/builder.py (1)
get_tools(110-113)
src/nat/control_flow/sequential_executor.py (2)
src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
examples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py (3)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (2)
src/nat/builder/function.py (2)
Function(46-301)get_accessible_functions(485-521)src/nat/builder/workflow_builder.py (1)
build(217-332)
tests/nat/builder/test_builder.py (2)
src/nat/builder/workflow_builder.py (3)
build(217-332)set_workflow(499-508)set_workflow(1157-1158)src/nat/builder/function.py (4)
get_accessible_functions(485-521)get_all_functions(628-664)get_included_functions(580-626)get_excluded_functions(523-578)
🪛 Ruff (0.13.1)
src/nat/agent/tool_calling_agent/register.py
65-65: Avoid specifying long messages outside the exception class
(TRY003)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py
179-179: Unused method argument: tool_names
(ARG002)
179-179: Unused method argument: wrapper_type
(ARG002)
tests/nat/builder/test_function_group.py
149-149: Unused function argument: name
(ARG001)
152-152: Unused function argument: name
(ARG001)
236-236: Unused function argument: name
(ARG001)
302-302: Unused function argument: name
(ARG001)
363-363: Unused function argument: names
(ARG001)
616-616: Unused function argument: name
(ARG001)
623-623: Unused function argument: name
(ARG001)
src/nat/builder/workflow_builder.py
454-454: Avoid specifying long messages outside the exception class
(TRY003)
src/nat/builder/function.py
458-458: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI Pipeline / Check
🔇 Additional comments (64)
examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.py (1)
75-76: Awaiting async get_tools: LGTMCorrectly updated to await the async API.
examples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.py (1)
78-79: Awaiting async get_tools: LGTMMatches the new async API surface.
docs/source/workflows/function-groups.md (2)
187-189: Docs: reflect async group filters: LGTMUsing await with get_accessible_functions is correct.
200-204: Docs: async get_tools usage: LGTMExamples now properly await builder.get_tools.
src/nat/agent/react_agent/register.py (1)
102-103: Awaiting async get_tools: LGTMCall site updated correctly.
src/nat/agent/rewoo_agent/register.py (1)
111-113: Awaiting async get_tools: LGTMConsistent with async API changes.
examples/notebooks/1_getting_started.ipynb (1)
257-259: Async get_tools in notebook: LGTMTop-level await is fine in IPython; ensure the kernel supports it.
Please confirm you’re on IPython 8+ where top-level
awaitis enabled.examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.py (1)
62-63: Awaiting async get_tools: LGTMMatches the new async contract.
examples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.py (2)
70-70: Async migration alignment: get_tool remains sync — LGTMUsing the synchronous
builder.get_tool(...)is correct givenget_toolwasn’t made async in this PR.
335-339: Verify LangChain tool method:arunvsainvokeMany LangChain tools expose
ainvoke(preferred) rather thanarun. Please confirm the wrapper forgit_repo_toolimplementsarun; if not, switch toawait self.git_tool.ainvoke(...).Would you like me to scan for remaining
.arunusages and propose a patch?src/nat/control_flow/router_agent/register.py (1)
56-56: Correct: await asyncget_toolsAwaiting
builder.get_tools(...)matches the new async API.src/nat/control_flow/sequential_executor.py (1)
113-114: Correct: await asyncget_toolsThe awaited call matches the updated API and surrounding async flow.
src/nat/front_ends/simple_base/simple_front_end_plugin_base.py (1)
48-48: Correct: await asyncbuilder.build()This aligns with the new async
WorkflowBuilder.build.examples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.py (1)
51-51: Correct: await asyncget_toolsMatches the async migration; downstream usage is unchanged.
tests/nat/front_ends/mcp/test_mcp_custom_routes.py (1)
67-67: Async mocks added — LGTMImporting
AsyncMockis required for the new async paths.examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.py (1)
61-61: Correct: await asyncget_toolsConsistent with the updated builder API.
src/nat/runtime/loader.py (1)
117-117: Correct: await async build before creating SessionManagerEnsures a realized
Workflowis passed toSessionManager.examples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py (1)
60-63: LGTM: awaited get_tools matches new async API.src/nat/eval/evaluate.py (1)
523-525: Correct: await build() before constructing SessionManager.tests/nat/front_ends/mcp/test_mcp_front_end_plugin.py (1)
57-64: LGTM: async tests and awaits reflect new async FunctionGroup API.Also applies to: 82-99, 112-121, 150-168, 178-205
src/nat/builder/builder.py (1)
110-113: Approve: async API conversion verified.
All get_tools implementations in the repo are async; no non-async overrides found.tests/nat/builder/test_function_group.py (18)
91-92: LGTM! Test correctly updated for async filter functions.The test decorator and function signature properly reflect the async nature of the filter callback.
97-98: LGTM! Filter function properly converted to async.The group filter function correctly follows the new async signature
Callable[[Sequence[str]], Awaitable[Sequence[str]]].
117-119: LGTM! All function group accessor methods correctly awaited.All calls to
get_accessible_functions(),get_all_functions(), and the filter override are properly awaited to match the new async API.Also applies to: 122-123, 126-126, 129-130
133-134: LGTM! Test decorator properly applied for async test.
149-153: Static analysis false positive: Unused parameter is intentional.The
nameparameter is part of the per-function filter signature contract but not needed for these specific test filters. This is correct test implementation.
160-161: LGTM! Function group accessor methods correctly awaited.Also applies to: 165-166
169-170: LGTM! Filter functions and accessor calls properly updated to async.The test correctly demonstrates filter interaction with include configuration in the async context.
Also applies to: 175-176, 189-190, 193-194
197-198: LGTM! Filter and accessor methods correctly updated for async.Also applies to: 203-204, 217-218
221-222: LGTM! Complex filter interaction test properly updated.The test correctly demonstrates the interaction between group filters, per-function filters, and configuration in an async context.
Also applies to: 227-228, 236-237, 246-247, 251-252
255-256: LGTM! Set filter test correctly updated for async.Also applies to: 270-270, 274-275, 280-281
284-285: LGTM! Per-function filter setter test correctly updated.Also applies to: 298-299, 302-303, 308-309
316-317: LGTM! Validation error tests correctly updated for async.Also applies to: 330-330, 333-333
336-337: LGTM! Exclude config validation tests correctly updated.Also applies to: 351-351, 354-354
357-358: LGTM! Empty filter behavior test correctly updated.Also applies to: 363-364, 375-376, 378-379
382-383: LGTM! Filter override precedence test correctly updated.Also applies to: 388-389, 400-401, 404-405, 407-408, 411-412
475-476: LGTM! Function name generation test correctly updated.Also applies to: 491-492
513-514: LGTM! Empty include/exclude test correctly updated.Also applies to: 531-532
535-536: LGTM! All remaining async test updates are correct.All test methods properly await the async function group methods and filter results.
Also applies to: 569-570, 581-582, 585-586, 598-599, 602-603, 613-613, 616-617, 620-620, 623-624, 627-627, 630-631, 667-668
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (3)
89-89: LGTM! Method signature correctly updated to async.The
_get_all_functionsmethod signature properly reflects the async nature required for accessing function groups.
103-103: LGTM! Function group accessor correctly awaited.The call to
get_accessible_functions()is properly awaited to match the new async API.
226-226: LGTM! Build and helper methods correctly awaited.Both
builder.build()and_get_all_functions()are properly awaited in the MCP worker implementation.Also applies to: 229-229
tests/nat/builder/test_builder.py (3)
427-427: LGTM! Builder.build() correctly awaited in tests.All calls to
builder.build()are properly awaited to match the new async API.Also applies to: 433-433, 921-921
1040-1040: LGTM! Function group accessor methods correctly awaited.All calls to function group methods (
get_accessible_functions(),get_all_functions(),get_included_functions(),get_excluded_functions()) are properly awaited.Also applies to: 1062-1065, 1085-1088, 1107-1109
1227-1227: LGTM! All remaining async method calls correctly awaited.Also applies to: 1237-1237, 1287-1287, 1291-1291, 1338-1338, 1342-1342, 1363-1363
docs/source/extend/function-groups.md (2)
93-93: LGTM! Documentation correctly updated for async API.The code examples properly show the
awaitkeyword for all async function group methods.Also applies to: 97-97, 100-100, 103-103
143-143: LGTM! Filter examples correctly updated to async.All filter function examples properly demonstrate the async signatures with
async defand appropriate return type annotations.Also applies to: 172-172, 205-205, 211-211
src/nat/builder/workflow_builder.py (6)
217-217: LGTM! Build method signature correctly updated to async.The
buildmethod is properly converted to async, matching the PR's breaking change requirements.
253-253: LGTM! Function group included functions correctly awaited.The call to
get_included_functions()is properly awaited when collecting included functions from function groups.
451-457: LGTM! Function group registration correctly handles async included functions.The code properly awaits
get_included_functions()and correctly iterates over the result.
538-540: LGTM! Get_tools method correctly updated to async.The method signature properly reflects the async requirement for tool retrieval.
563-563: LGTM! Function group accessible functions correctly awaited.The call to
get_accessible_functions()is properly awaited when iterating through function group functions.
1169-1172: LGTM! ChildBuilder get_tools correctly updated to async.The method properly awaits the parent's
get_toolscall and maintains the async contract.src/nat/builder/function.py (11)
360-373: LGTM! Proper async type annotations for filter functions.The function signature correctly adopts async filter functions, using proper type hints with
Awaitable[Sequence[str]]as described in the PR objectives. The docstring is also updated to reflect the new async behavior.
378-378: Correct update to per-function filter storage.The type annotation properly reflects the new async requirement for per-function filters, storing
Callable[[str], Awaitable[bool]]callbacks.
387-407: LGTM! Function signature and docstring properly updated.The
add_functionmethod correctly adopts the new async filter function signatureCallable[[str], Awaitable[bool]]and the docstring accurately describes the new async behavior.
444-447: LGTM! Proper async implementation for per-function filtering.The method is correctly converted to async and properly awaits the per-function filter callback result. The logic remains correct - returning
Truewhen no filter exists and awaiting the filter result when one is present.
449-483: LGTM! Robust async implementation with proper error handling.The method correctly:
- Uses proper async type annotations
- Implements async identity filter as fallback
- Awaits both the main filter function and per-function filters
- Maintains proper validation of excluded function names
- Uses efficient set operations for filtering logic
485-521: LGTM! Proper async delegation to helper methods.The method correctly becomes async and delegates to the appropriate async helper methods (
get_included_functions,_get_all_but_excluded_functions,get_all_functions) based on configuration, maintaining the same logical flow.
523-578: LGTM! Consistent async implementation for excluded functions.The method properly:
- Adopts async signature with correct type annotations
- Implements the same async identity filter pattern
- Awaits filter results appropriately
- Maintains validation logic for excluded functions
- Uses consistent variable naming and logic flow
580-626: LGTM! Proper async implementation for included functions.The method correctly:
- Uses async signature with appropriate type hints
- Implements consistent async identity filter fallback
- Awaits filter function results
- Validates included function names against actual functions
- Properly awaits per-function filter checks
628-664: LGTM! Clean async implementation for all functions.The method properly implements async patterns consistently with other methods, using the same identity filter fallback pattern and awaiting both main filter and per-function filter results.
666-675: LGTM! Method signature correctly updated.The
set_filter_fnmethod signature is properly updated to accept the new async filter function typeCallable[[Sequence[str]], Awaitable[Sequence[str]]].
677-695: LGTM! Per-function filter setter properly updated.The
set_per_function_filter_fnmethod signature is correctly updated to acceptCallable[[str], Awaitable[bool]]and maintains proper validation for function existence.
willkill07
force-pushed
the
wkk_async-function-group-callback
branch
from
ab5a6d4 to
716f694
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/nat/front_ends/simple_base/simple_front_end_plugin_base.py (1)
31-54: Public API lacks required type hints and Google‑style docstrings.Methods
pre_run,run, and abstractrun_workfloware public and should have explicit return types and docstrings per repo guidelines.Proposed update:
class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC): """Base plugin for simple front ends. Orchestrates building a `Workflow` from `full_config` and delegates execution to `run_workflow`. """ async def pre_run(self) -> None: """Optional hook invoked before running the workflow.""" pass async def run(self) -> None: """Build the workflow from `full_config` and execute `run_workflow`.""" ... @abstractmethod async def run_workflow(self, session_manager: SessionManager) -> None: """Run the workflow loop. Args: session_manager: The session manager bound to the built workflow. """ passsrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (1)
89-111: Prevent workflow-name collisions; optionally parallelize function-group fetch.
- Risk: adding the workflow under
workflow_aliasor fallbackworkflow.typecan overwrite an existing function of the same name.- Fix: guard and fail fast with a clear message. Optionally parallelize
get_accessible_functions()across groups to reduce latency.@@ - for function_group in workflow.function_groups.values(): - functions.update(await function_group.get_accessible_functions()) + # Collect group functions (optionally parallelized below) + for function_group in workflow.function_groups.values(): + functions.update(await function_group.get_accessible_functions()) @@ - if workflow.config.workflow.workflow_alias: - functions[workflow.config.workflow.workflow_alias] = workflow - else: - functions[workflow.config.workflow.type] = workflow + wf_name = (workflow.config.workflow.workflow_alias + or workflow.config.workflow.type) + if wf_name in functions: + raise ValueError( + f"Workflow tool name `{wf_name}` collides with an existing function. " + "Please set a unique `workflow_alias` in the workflow config." + ) + functions[wf_name] = workflowOptional parallelization inside the selected range (requires
import asyncioat file top):@@ - for function_group in workflow.function_groups.values(): - functions.update(await function_group.get_accessible_functions()) + group_funcs_list = await asyncio.gather( + *(fg.get_accessible_functions() for fg in workflow.function_groups.values()) + ) + for gf in group_funcs_list: + functions.update(gf)Add once at the top of the file:
import asynciosrc/nat/builder/workflow_builder.py (1)
1169-1179: Fix dependency tracking for FunctionRef/FunctionGroupRef.
tool_namemay be a ref object; membership checks against dicts keyed bystrwill fail, causing missing dependency edges.- for tool_name in tool_names: - if tool_name in self._workflow_builder._function_groups: - self._dependencies.add_function_group(tool_name) - else: - self._dependencies.add_function(tool_name) + for tool_name in tool_names: + n = (str(tool_name) + if isinstance(tool_name, (FunctionRef, FunctionGroupRef)) + else tool_name) + if n in self._workflow_builder._function_groups: + self._dependencies.add_function_group(n) + else: + self._dependencies.add_function(n)
🧹 Nitpick comments (15)
src/nat/front_ends/simple_base/simple_front_end_plugin_base.py (1)
36-50: Consider invoking thepre_runhook.If nothing upstream calls
pre_run, invoke it here so subclasses can initialize before build/execute.async def run(self): # Must yield the workflow function otherwise it cleans up async with WorkflowBuilder.from_config(config=self.full_config) as builder: if logger.isEnabledFor(logging.INFO): stream = StringIO() self.full_config.print_summary(stream=stream) click.echo(stream.getvalue()) + await self.pre_run() workflow = await builder.build() session_manager = SessionManager(workflow) await self.run_workflow(session_manager)If
pre_runis already called by the framework, ignore this change.tests/nat/builder/test_component_utils.py (1)
436-436: Strengthen the assertion for clarity.Assert the constructed type explicitly.
- assert SessionManager(await workflow.build(), max_concurrency=1) + assert isinstance(SessionManager(await workflow.build(), max_concurrency=1), SessionManager)examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (1)
82-89: Remove duplicate tool resolution; use a single awaited call.The manual loop builds tools, then you reassign with awaited get_tools. Keep only the awaited path to avoid divergence and double work.
@@ - # Get tools for alert triage - tool_names = config.tool_names - tools = [] - for tool_name in tool_names: - tool = builder.get_tool(tool_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN) - tools.append(tool) - llm_n_tools = llm.bind_tools(tools, parallel_tool_calls=True) + # Get tools for alert triage + tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + llm_n_tools = llm.bind_tools(tools, parallel_tool_calls=True) @@ - tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + # tools already resolved aboveAlso applies to: 104-104
src/nat/agent/tool_calling_agent/register.py (1)
63-70: Trim exception message to satisfy ruff TRY003 (and add context via logging).Shorten the ValueError message and log details.
tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) if not tools: - raise ValueError(f"No tools specified for Tool Calling Agent '{config.llm_name}'") + logger.error("No tools specified for Tool Calling Agent %s", config.llm_name) + raise ValueError("No tools specified") @@ - return_direct_tools = await builder.get_tools( - tool_names=config.return_direct, wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None + return_direct_tools = ( + await builder.get_tools(tool_names=config.return_direct, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + if config.return_direct else None + )src/nat/builder/workflow_builder.py (1)
217-333: Async build LGTM; consider parallelizing function-group includes.Current per-group awaits are sequential. You can shave build latency by gathering included functions concurrently.
@@ - for k, v in self._function_groups.items(): - included_functions.update((await v.instance.get_included_functions()).keys()) - function_group_configs[k] = v.config - function_group_instances[k] = v.instance + # Collect configs/instances first + for k, v in self._function_groups.items(): + function_group_configs[k] = v.config + function_group_instances[k] = v.instance + # Fetch includes concurrently + include_maps = await asyncio.gather( + *(v.instance.get_included_functions() for v in self._function_groups.values()) + ) + for inc in include_maps: + included_functions.update(inc.keys())Note: add
import asyncioat the top of this module if you adopt this.src/nat/builder/function.py (5)
444-448: Per-function async filter LGTM.Tiny nit: use
fn = self._per_function_filter_fn.get(name)to avoid double dict lookups (micro).
449-484: Harden filter handling; avoid KeyError if filter returns unknown names.Intersect filtered results with known keys; dedupe repeated identity-filter logic later (DRY).
@@ - if filter_fn is None: - if self._filter_fn is None: - async def identity_filter(x: Sequence[str]) -> Sequence[str]: - return x - filter_fn = identity_filter - else: - filter_fn = self._filter_fn + if filter_fn is None: + if self._filter_fn is None: + async def identity_filter(x: Sequence[str]) -> Sequence[str]: + return x + filter_fn = identity_filter + else: + filter_fn = self._filter_fn @@ - included = set(await filter_fn(list(self._functions.keys()))) + included = set(await filter_fn(list(self._functions.keys()))) + included &= set(self._functions.keys()) @@ - for name in self._functions: + for name in self._functions: if name in excluded: continue if not await self._fn_should_be_included(name): continue if name not in included: continue result[self._get_fn_name(name)] = self._functions[name]Optional: introduce a private helper like
_resolve_seq_filter(filter_fn)returning an async callable to remove the repeated identity-filter blocks.
523-579: Harden excluded-functions path similarly.Guard against filters returning unknown names; keeps behavior robust.
@@ - if filter_fn is None: - if self._filter_fn is None: - async def identity_filter(x: Sequence[str]) -> Sequence[str]: - return x - filter_fn = identity_filter - else: - filter_fn = self._filter_fn + if filter_fn is None: + if self._filter_fn is None: + async def identity_filter(x: Sequence[str]) -> Sequence[str]: + return x + filter_fn = identity_filter + else: + filter_fn = self._filter_fn @@ - included = set(await filter_fn(list(self._functions.keys()))) + included = set(await filter_fn(list(self._functions.keys()))) + included &= set(self._functions.keys())
580-627: Harden included-functions path; minor DRY opportunity.Intersect filtered names with known functions to avoid accidental KeyError if a filter misbehaves.
@@ - if filter_fn is None: - if self._filter_fn is None: - async def identity_filter(x: Sequence[str]) -> Sequence[str]: - return x - filter_fn = identity_filter - else: - filter_fn = self._filter_fn + if filter_fn is None: + if self._filter_fn is None: + async def identity_filter(x: Sequence[str]) -> Sequence[str]: + return x + filter_fn = identity_filter + else: + filter_fn = self._filter_fn @@ - included = set(await filter_fn(list(self._config.include))) + included = set(await filter_fn(list(self._config.include))) + included &= set(self._functions.keys())
628-665: Harden all-functions path; DRY the identity-filter.Same intersection + shared helper suggestion.
@@ - if filter_fn is None: - if self._filter_fn is None: - async def identity_filter(x: Sequence[str]) -> Sequence[str]: - return x - filter_fn = identity_filter - else: - filter_fn = self._filter_fn + if filter_fn is None: + if self._filter_fn is None: + async def identity_filter(x: Sequence[str]) -> Sequence[str]: + return x + filter_fn = identity_filter + else: + filter_fn = self._filter_fn @@ - included = set(await filter_fn(list(self._functions.keys()))) + included = set(await filter_fn(list(self._functions.keys()))) + included &= set(self._functions.keys())Optional helper (outside selected ranges):
@staticmethod async def _identity_seq_filter(x: Sequence[str]) -> Sequence[str]: return x def _resolve_seq_filter( self, filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None ) -> Callable[[Sequence[str]], Awaitable[Sequence[str]]]: return filter_fn or self._filter_fn or self._identity_seq_filtertests/nat/builder/test_function_group.py (5)
149-154: Fix ruff ARG001: unused parameter in per-function filters.Rename the unused parameter to underscore to silence ARG001.
- async def exclude_func1(name: str) -> bool: + async def exclude_func1(_name: str) -> bool: return False # Always exclude func1 - async def include_func2(name: str) -> bool: + async def include_func2(_name: str) -> bool: return True # Always include func2
236-236: Fix ruff ARG001: unused parameter.Rename the unused parameter in exclude_func2.
- async def exclude_func2(name: str) -> bool: + async def exclude_func2(_name: str) -> bool: return False
302-302: Fix ruff ARG001: unused parameter.Rename the unused parameter in exclude_func1.
- async def exclude_func1(name: str) -> bool: + async def exclude_func1(_name: str) -> bool: return False
363-363: Fix ruff ARG001: unused parameter in empty filter.Rename the parameter to underscore.
- async def empty_filter(names: Sequence[str]) -> Sequence[str]: + async def empty_filter(_names: Sequence[str]) -> Sequence[str]: return []
616-616: Fix ruff ARG001: unused parameters in per-function filters.Rename the unused parameters.
- async def exclude_filter(name: str) -> bool: + async def exclude_filter(_name: str) -> bool: return False @@ - async def include_filter(name: str) -> bool: + async def include_filter(_name: str) -> bool: return TrueAlso applies to: 623-623
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (31)
docs/source/extend/function-groups.md(4 hunks)docs/source/workflows/function-groups.md(2 hunks)examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py(1 hunks)examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.py(1 hunks)examples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.py(1 hunks)examples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.py(1 hunks)examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.py(1 hunks)examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.py(1 hunks)examples/notebooks/1_getting_started.ipynb(1 hunks)examples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.py(1 hunks)examples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py(1 hunks)packages/nvidia_nat_test/src/nat/test/tool_test_runner.py(1 hunks)src/nat/agent/react_agent/register.py(1 hunks)src/nat/agent/rewoo_agent/register.py(1 hunks)src/nat/agent/tool_calling_agent/register.py(1 hunks)src/nat/builder/builder.py(1 hunks)src/nat/builder/function.py(12 hunks)src/nat/builder/workflow_builder.py(6 hunks)src/nat/control_flow/router_agent/register.py(1 hunks)src/nat/control_flow/sequential_executor.py(1 hunks)src/nat/eval/evaluate.py(1 hunks)src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py(1 hunks)src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py(3 hunks)src/nat/front_ends/simple_base/simple_front_end_plugin_base.py(1 hunks)src/nat/runtime/loader.py(1 hunks)tests/nat/builder/test_builder.py(11 hunks)tests/nat/builder/test_component_utils.py(1 hunks)tests/nat/builder/test_function_group.py(22 hunks)tests/nat/front_ends/mcp/test_mcp_custom_routes.py(1 hunks)tests/nat/front_ends/mcp/test_mcp_front_end_plugin.py(8 hunks)tests/nat/mcp/test_mcp_client_impl.py(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (19)
- src/nat/agent/rewoo_agent/register.py
- tests/nat/front_ends/mcp/test_mcp_custom_routes.py
- src/nat/runtime/loader.py
- examples/frameworks/agno_personal_finance/src/nat_agno_personal_finance/agno_personal_finance_function.py
- examples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.py
- src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py
- src/nat/control_flow/router_agent/register.py
- src/nat/agent/react_agent/register.py
- examples/advanced_agents/profiler_agent/src/nat_profiler_agent/register.py
- docs/source/workflows/function-groups.md
- src/nat/builder/builder.py
- examples/notebooks/first_search_agent/src/nat_first_search_agent/second_search_agent_function.py
- src/nat/control_flow/sequential_executor.py
- examples/notebooks/1_getting_started.ipynb
- src/nat/eval/evaluate.py
- docs/source/extend/function-groups.md
- examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/register.py
- examples/notebooks/retail_sales_agent/src/nat_retail_sales_agent/data_visualization_agent.py
- examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/telemetry_metrics_analysis_agent.py
🧰 Additional context used
📓 Path-based instructions (12)
**/*.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:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pysrc/nat/agent/tool_calling_agent/register.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:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pysrc/nat/agent/tool_calling_agent/register.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:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pysrc/nat/agent/tool_calling_agent/register.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:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pysrc/nat/agent/tool_calling_agent/register.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:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.pytests/nat/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pypackages/nvidia_nat_test/src/nat/test/tool_test_runner.pysrc/nat/agent/tool_calling_agent/register.py
examples/**/*
⚙️ CodeRabbit configuration file
examples/**/*: - This directory contains example code and usage scenarios for the toolkit, at a minimum an example should
contain a README.md or file README.ipynb.
- If an example contains Python code, it should be placed in a subdirectory named
src/and should
contain apyproject.tomlfile. Optionally, it might also contain scripts in ascripts/directory.- If an example contains YAML files, they should be placed in a subdirectory named
configs/. - If an example contains sample data files, they should be placed in a subdirectory nameddata/, and should
be checked into git-lfs.
Files:
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.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/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.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/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.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/front_ends/mcp/test_mcp_front_end_plugin.pytests/nat/builder/test_function_group.pytests/nat/mcp/test_mcp_client_impl.pytests/nat/builder/test_component_utils.pytests/nat/builder/test_builder.py
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/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pysrc/nat/agent/tool_calling_agent/register.py
src/nat/**/*
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Core functionality under src/nat should prioritize backward compatibility when changed
Files:
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pysrc/nat/agent/tool_calling_agent/register.py
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/simple_base/simple_front_end_plugin_base.pysrc/nat/builder/function.pysrc/nat/builder/workflow_builder.pysrc/nat/agent/tool_calling_agent/register.py
packages/*/src/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
packages/*/src/**/*.py: All importable Python code in packages must live under packages//src/
All public APIs in packaged code require Python 3.11+ type hints; prefer typing/collections.abc; use typing.Annotated when useful
Provide Google-style docstrings for public APIs in packages; first line concise with a period; use backticks for code entities
Files:
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py
packages/**/*
⚙️ CodeRabbit configuration file
packages/**/*: - This directory contains optional plugin packages for the toolkit, each should contain apyproject.tomlfile. - Thepyproject.tomlfile should declare a dependency onnvidia-nator another package with a name starting
withnvidia-nat-. This dependency should be declared using~=<version>, and the version should be a two
digit version (ex:~=1.0).
- Not all packages contain Python code, if they do they should also contain their own set of tests, in a
tests/directory at the same level as thepyproject.tomlfile.
Files:
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py
🧠 Learnings (1)
📚 Learning: 2025-09-09T20:32:39.016Z
Learnt from: CR
PR: NVIDIA/NeMo-Agent-Toolkit#0
File: .cursor/rules/nat-test-llm.mdc:0-0
Timestamp: 2025-09-09T20:32:39.016Z
Learning: Applies to **/*.py : When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.<FRAMEWORK>) and call the framework’s method (e.g., ainvoke, achat, call).
Applied to files:
src/nat/agent/tool_calling_agent/register.py
🧬 Code graph analysis (11)
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (2)
src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
tests/nat/front_ends/mcp/test_mcp_front_end_plugin.py (2)
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (1)
_get_all_functions(89-110)src/nat/builder/function.py (1)
get_accessible_functions(485-521)
tests/nat/builder/test_function_group.py (1)
src/nat/builder/function.py (8)
get_accessible_functions(485-521)get_all_functions(628-664)add_function(380-428)get_included_functions(580-626)get_excluded_functions(523-578)set_filter_fn(666-675)set_per_function_filter_fn(677-695)_fn_should_be_included(444-447)
tests/nat/mcp/test_mcp_client_impl.py (1)
src/nat/builder/function.py (1)
get_accessible_functions(485-521)
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (3)
src/nat/builder/workflow.py (1)
Workflow(41-160)src/nat/builder/function.py (2)
Function(46-301)get_accessible_functions(485-521)src/nat/builder/workflow_builder.py (1)
build(217-332)
src/nat/front_ends/simple_base/simple_front_end_plugin_base.py (2)
examples/object_store/user_report/tests/test_user_report_tool.py (1)
builder(33-49)src/nat/builder/workflow_builder.py (1)
build(217-332)
tests/nat/builder/test_component_utils.py (2)
src/nat/runtime/session.py (2)
SessionManager(45-173)workflow(84-85)src/nat/builder/workflow_builder.py (1)
build(217-332)
tests/nat/builder/test_builder.py (2)
src/nat/builder/workflow_builder.py (3)
build(217-332)set_workflow(499-508)set_workflow(1157-1158)src/nat/builder/function.py (4)
get_accessible_functions(485-521)get_all_functions(628-664)get_included_functions(580-626)get_excluded_functions(523-578)
src/nat/builder/workflow_builder.py (2)
src/nat/builder/function.py (2)
get_included_functions(580-626)get_accessible_functions(485-521)src/nat/builder/builder.py (1)
get_tools(110-113)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (2)
src/nat/builder/builder.py (1)
get_tools(110-113)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
src/nat/agent/tool_calling_agent/register.py (2)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (1)
get_tools(179-181)src/nat/builder/workflow_builder.py (2)
get_tools(538-571)get_tools(1169-1178)
🪛 Ruff (0.13.1)
tests/nat/builder/test_function_group.py
149-149: Unused function argument: name
(ARG001)
152-152: Unused function argument: name
(ARG001)
236-236: Unused function argument: name
(ARG001)
302-302: Unused function argument: name
(ARG001)
363-363: Unused function argument: names
(ARG001)
616-616: Unused function argument: name
(ARG001)
623-623: Unused function argument: name
(ARG001)
src/nat/builder/function.py
458-458: Avoid specifying long messages outside the exception class
(TRY003)
src/nat/builder/workflow_builder.py
454-454: Avoid specifying long messages outside the exception class
(TRY003)
packages/nvidia_nat_test/src/nat/test/tool_test_runner.py
179-179: Unused method argument: tool_names
(ARG002)
179-179: Unused method argument: wrapper_type
(ARG002)
src/nat/agent/tool_calling_agent/register.py
65-65: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (25)
src/nat/front_ends/simple_base/simple_front_end_plugin_base.py (1)
48-50: Async build usage verified — no unawaited build() call sites found.Found build() usages in src/nat/runtime/loader.py:117, src/nat/eval/evaluate.py:523, src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py:226, src/nat/front_ends/simple_base/simple_front_end_plugin_base.py:48, src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py:240,241,247 and in tests; each occurrence is awaited.
tests/nat/mcp/test_mcp_client_impl.py (1)
124-126: LGTM: awaiting async FunctionGroup accessors.The tests correctly await get_accessible_functions(), matching the async API.
Also applies to: 143-146, 159-161
tests/nat/builder/test_builder.py (2)
427-427: LGTM: build() is now awaited.Consistent with async WorkflowBuilder.build.
Also applies to: 433-433, 921-921, 1364-1364
1041-1041: LGTM: FunctionGroup queries now awaited.Verified with an AST scan — 83 target calls found, 0 non-awaited occurrences.
tests/nat/front_ends/mcp/test_mcp_front_end_plugin.py (1)
57-59: LGTM: async tests and awaited _get_all_functions.Tests correctly adopt asyncio style and await worker helpers; function-group AsyncMock usage is appropriate.
Also applies to: 73-79, 83-85, 97-99, 112-114, 133-139, 143-149, 152-156, 170-176, 178-185, 189-193, 204-213
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (1)
226-230: Async propagation LGTM.Awaiting
builder.build()and_get_all_functions()is correct and matches the new async contracts.src/nat/builder/workflow_builder.py (2)
451-458: Function-group exposure to global registry LGTM.Awaiting
get_included_functions()and conflict guarding before updating_functionsis correct.
538-571: Async get_tools LGTM; error handling is appropriate.Type coercion and logging before re-raise look good.
src/nat/builder/function.py (3)
360-379: Async filter signatures LGTM; document “subset” contract.Interfaces are correct. Please clarify in docstrings that filters must return a subset of provided names (not arbitrary names) to avoid undefined lookups.
Also applies to: 387-388, 403-404
485-522: Async access orchestration LGTM.Control flow for include/exclude/all is correct.
666-676: Setter APIs LGTM.Types and validation look good.
Also applies to: 677-696
tests/nat/builder/test_function_group.py (14)
91-131: Async migration in tests looks correct.Proper use of pytest.mark.asyncio and awaiting async accessors; override filter precedence is covered.
169-195: Include-config x async filter interaction is validated.The assertions match the new async contract.
197-219: Exclude-config x async filter interaction is validated.Happy path and intersection semantics look good.
246-253: Complex filter interactions covered well.Good coverage of excluded functions contract.
255-282: set_filter_fn async behavior is exercised.Confirms instance-level filter is applied post‑set.
316-334: Config validation errors for includes are asserted correctly.Awaiting the async accessors is correct.
337-355: Config validation errors for excludes are asserted correctly.Covers both direct and indirect accessor paths.
375-380: Empty filter behavior validated.Both accessible and all functions correctly yield empty.
382-413: Override precedence between parameter filter and instance filter is correct.Covers both override and fallback paths.
475-494: Function name generation with custom instance_name is verified.Good assertion on full vs. short names.
513-533: Empty include/exclude lists semantics are correct.Confirms default “all accessible” path.
569-583: No-exclusions path validated.get_excluded_functions returns empty as expected.
585-600: No-includes semantics validated.get_included_functions returns empty dict as config.include is empty.
630-670: Comprehensive metadata test LGTM.Covers converters, schemas, and instance_name; async accessor usage is correct.
zhongxuanwang-nv
left a comment
zhongxuanwang-nv
left a comment
There was a problem hiding this comment.
Thanks Will! I think this looks great!
willkill07
changed the title
async in Builder's get_* functions; make function group filters async
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (13)
examples/advanced_agents/alert_triage_agent/tests/test_telemetry_metrics_host_heartbeat_check_tool.py (1)
92-94: Patch async helper with AsyncMock to avoid await on MagicMock
llm_ainvokeis async; patch it withAsyncMockso awaiting it in the tool under test doesn’t raise TypeError.Apply this diff:
-with patch('requests.get') as mock_get, \ - patch('nat_alert_triage_agent.utils.llm_ainvoke') as mock_llm_invoke: +with patch('requests.get') as mock_get, \ + patch('nat_alert_triage_agent.utils.llm_ainvoke', new_callable=AsyncMock) as mock_llm_invoke:examples/frameworks/multi_frameworks/src/nat_multi_frameworks/langchain_research_tool.py (1)
45-51: Fix None assignment to environment variable (runtime error)
os.environ["NVIDIA_API_KEY"] = api_tokenwill crash whenapi_tokenis None. Guard or remove the assignment.- api_token = os.getenv("NVIDIA_API_KEY") - os.environ["NVIDIA_API_KEY"] = api_token - - if not api_token: + api_token = os.getenv("NVIDIA_API_KEY") + if not api_token: raise ValueError( "API token must be provided in the configuration or in the environment variable `NVIDIA_API_KEY`") + # No need to reset the env var if already presentexamples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.py (1)
258-264: Patch async function with AsyncMock
_summarize_alertis likely awaited by the tool; patch it withAsyncMockto avoid await errors.- with patch('nat_alert_triage_agent.maintenance_check._summarize_alert') as mock_summarize: + with patch('nat_alert_triage_agent.maintenance_check._summarize_alert', + new_callable=AsyncMock) as mock_summarize:examples/advanced_agents/alert_triage_agent/tests/test_telemetry_metrics_host_performance_check_tool.py (1)
106-108: Patch async helper with AsyncMock
llm_ainvokeshould be patched asAsyncMockto be awaitable.-with patch('requests.get') as mock_get, \ - patch('nat_alert_triage_agent.utils.llm_ainvoke') as mock_llm_invoke: +with patch('requests.get') as mock_get, \ + patch('nat_alert_triage_agent.utils.llm_ainvoke', new_callable=AsyncMock) as mock_llm_invoke:.cursor/rules/nat-workflows/add-functions.mdc (3)
199-224: Docs bug: yields wrong symbol (my_callableundefined).The example defines
_my_functionbut yieldsmy_callable. Yield the defined function.Apply this diff:
- yield my_callable + yield _my_function
331-339: Generic parameter order is incorrect.
Functionis parameterized as[Input, StreamingOutput, SingleOutput]. The example swaps single/streaming.Apply this diff:
-class MyFunction(Function[MyInput, MySingleOutput, MyStreamingOutput]): +class MyFunction(Function[MyInput, MyStreamingOutput, MySingleOutput]):
364-377: Stream-to-single converter must be async.Converting an
AsyncGeneratorrequiresasync defwithasync for. The current sync example won’t work.Apply this diff:
-# Define a conversion function to convert a streaming output to a single output -def convert_streaming_to_single(value: AsyncGenerator[str]) -> str: - return ", ".join(value) +# Define a conversion function to convert a streaming output to a single output +async def convert_streaming_to_single(value: AsyncGenerator[str, None]) -> str: + items: list[str] = [] + async for item in value: + items.append(item) + return ", ".join(items)docs/source/extend/functions.md (3)
199-224: Docs bug: yields wrong symbol (my_callableundefined).The example defines
_my_functionbut yieldsmy_callable. Yield the defined function.Apply this diff:
- yield my_callable + yield _my_function
331-339: Generic parameter order is incorrect.
Function[Input, StreamingOutput, SingleOutput]— the example swaps them.Apply this diff:
-class MyFunction(Function[MyInput, MySingleOutput, MyStreamingOutput]): +class MyFunction(Function[MyInput, MyStreamingOutput, MySingleOutput]):
364-377: Stream-to-single converter must be async.Converting an
AsyncGeneratorrequiresasync defandasync for.Apply this diff:
-def convert_streaming_to_single(value: AsyncGenerator[str]) -> str: - return ", ".join(value) +async def convert_streaming_to_single(value: AsyncGenerator[str, None]) -> str: + items: list[str] = [] + async for item in value: + items.append(item) + return ", ".join(items)tests/nat/control_flow/test_sequential_executor.py (2)
565-568: Make side_effect async for missing tool fixture.Since get_tools is awaited, side_effect must be on an AsyncMock.
Apply:
- builder = MagicMock(spec=Builder) - builder.get_tools.side_effect = KeyError("Tool not found") + builder = MagicMock(spec=Builder) + builder.get_tools = AsyncMock(side_effect=KeyError("Tool not found"))
586-596: Use AsyncMock for get_tools and get_function in this test.These are awaited by production code; MagicMock will fail under await.
Apply:
- builder = MagicMock(spec=Builder) + builder = MagicMock(spec=Builder) tool = MockTool(name="tool1") - builder.get_tools.return_value = [tool] + builder.get_tools = AsyncMock(return_value=[tool]) @@ - builder.get_function.return_value = func + builder.get_function = AsyncMock(return_value=func)src/nat/control_flow/sequential_executor.py (1)
125-126: Preserve stack trace when re-raising.Wrap-and-raise loses the original stack. Log the error and use a bare
raiseper project guidelines.Apply:
- except Exception as e: - raise ValueError(f"Error with the sequential executor tool list: {e}") + except Exception: + logger.error("Error with the sequential executor tool list", exc_info=True) + raise
🧹 Nitpick comments (31)
src/nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py (1)
86-91: Duplicate output check; remove the redundant block.
has_single_outputis checked twice; keep one.Apply:
- if not augmented_function.has_single_output: - raise ValueError("TTCToolWrapperFunction only supports functions with a single output.") - - if not augmented_function.has_single_output: - raise ValueError("TTCToolWrapperFunction only supports functions with a single output.") + if not augmented_function.has_single_output: + raise ValueError("TTCToolWrapperFunction only supports functions with a single output.")examples/advanced_agents/alert_triage_agent/tests/test_telemetry_metrics_host_heartbeat_check_tool.py (1)
33-33: Mark async test with pytest-asyncio (if not globally configured)Add
@pytest.mark.asynciounless your pytest config enables asyncio auto mode.+@pytest.mark.asyncio async def test_telemetry_metrics_host_heartbeat_check_tool():examples/frameworks/multi_frameworks/src/nat_multi_frameworks/langchain_research_tool.py (3)
55-60: Correct the return type annotation for web_searchFunction returns a string (first segment), not
list[dict].- async def web_search(topic: str) -> list[dict]: + async def web_search(topic: str) -> str: output = (await tavily_tool.ainvoke(topic)) output = output.split("\n\n---\n\n") return output[0]
84-101: Tighten types and fix user-facing grammar
- Type the
outparam to the structured model.- Fix message grammar (“yielded no results”).
- async def execute_tool(out): + async def execute_tool(out: TopicExtract) -> str: try: topic = out.topic if topic is not None and topic not in ['', '\n']: output_summary = (await web_search(topic)) # Clean HTML tags from the output if isinstance(output_summary, str): # Remove HTML tags using BeautifulSoup soup = BeautifulSoup(output_summary, 'html.parser') output_summary = soup.get_text() # Clean up any extra whitespace output_summary = re.sub(r'\s+', ' ', output_summary).strip() else: - output_summary = f"this search on web search with topic:{topic} yield not results" + output_summary = f"this search on web search with topic:{topic} yielded no results" except Exception as e: - output_summary = f"this search on web search with topic:{topic} yield not results with an error:{e}" + output_summary = f"this search on web search with topic:{topic} yielded no results due to error: {e}" logger.exception("error in executing tool: %s", e) - pass
119-119: Typo in description: “relevant”Fix spelling in user-facing description.
- yield FunctionInfo.from_fn(_arun, description="extract relevent information from search the web") + yield FunctionInfo.from_fn(_arun, description="extract relevant information from searching the web")examples/HITL/por_to_jiratickets/src/nat_por_to_jiratickets/jira_tickets_tool.py (2)
346-352: Avoid blocking requests in async function
requests.getblocks the event loop inside async_arun. Preferhttpx.AsyncClient.- async def _arun(input_text: str) -> str: - response = requests.get(api_endpoint, headers=headers, params=query_params, timeout=30) + async def _arun(input_text: str) -> str: + async with httpx.AsyncClient(timeout=30) as client: + response = await client.get(api_endpoint, headers=headers, params=query_params)
68-71: Return type annotation mismatch for create_epicFunction returns a tuple on success and a dict on error, but the annotation says
-> str. Clarify the union type.- async def create_epic(self, client: httpx.AsyncClient, ticket_data: dict) -> str: + async def create_epic(self, client: httpx.AsyncClient, ticket_data: dict) -> tuple[str, str] | dict[str, str]:(Apply similar annotations to
create_task,create_bug, andcreate_featurefor consistency.)examples/HITL/simple_calculator_hitl/src/nat_simple_calculator_hitl/retry_react_agent.py (2)
167-167: Ensure return type matches ChatResponse
responsemay be a string when retries are exhausted. Wrap withChatResponse.from_stringfor type safety.- return response + return response if isinstance(response, ChatResponse) else ChatResponse.from_string(str(response))
207-209: Log unexpected exceptions with stack tracePer guidelines, use
logger.exceptionwhen catching and not re-raising.- except Exception: - # Handle any other unexpected exceptions - return ChatResponse.from_string("I seem to be having a problem.") + except Exception as e: + logger.exception("Unexpected error in retry_react_agent: %s", e) + return ChatResponse.from_string("I seem to be having a problem.")examples/advanced_agents/alert_triage_agent/tests/test_maintenance_check.py (1)
180-180: Mark async test with pytest-asyncio (if not auto)Add the asyncio marker unless globally configured.
+@pytest.mark.asyncio async def test_maintenance_check_tool():examples/advanced_agents/alert_triage_agent/tests/test_telemetry_metrics_host_performance_check_tool.py (1)
37-37: Mark async test with pytest-asyncio (if not auto)Add
@pytest.mark.asynciounless configured globally.+@pytest.mark.asyncio async def test_telemetry_metrics_host_performance_check_tool():examples/evaluation_and_profiling/swe_bench/src/nat_swe_bench/predictors/predict_full/predict_full.py (1)
26-26: Avoid blocking the event loop: use AsyncOpenAI.
OpenAI.chat.completions.create(...)is synchronous; calling it insideasync defcan block. PreferAsyncOpenAI.Apply this diff:
-from openai import OpenAI +from openai import AsyncOpenAI @@ - self.git_tool = None - self.openai_client = OpenAI(api_key=config.predictor.openai_api_key) + self.git_tool = None + self.openai_client = AsyncOpenAI(api_key=config.predictor.openai_api_key) @@ - logger.info("Generating fix using OpenAI GPT-4") - response = self.openai_client.chat.completions.create( + logger.info("Generating fix using OpenAI GPT-4") + response = await self.openai_client.chat.completions.create( model="gpt-4o", messages=[{ "role": "system", "content": "You are a skilled Python developer. Do not refactor code unnecessarily. Also, when printing output, print only the changes you made and not unchanged code. Generate only the fixed code without any explanations or formatting. " }, { "role": "user", "content": prompt }], temperature=0.2, max_tokens=2000)Also applies to: 62-66, 301-317
examples/frameworks/multi_frameworks/src/nat_multi_frameworks/register.py (1)
86-95: Avoid redundant tool fetching.You fetch tools via
_get_tool/gather, then refetch withbuilder.get_tools(...). Remove the second fetch and reuse the first list.Apply this diff:
- tools = await asyncio.gather(*tools) + tools = await asyncio.gather(*tools) @@ - tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + # Reuse the already-fetched tools for the graphAlso applies to: 107-107
examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (1)
89-91: Remove duplicate tool retrieval.You fetch tools via
gather, then fetch again withbuilder.get_tools(...). Drop the second call and reuse the firsttoolslist forToolNode.Apply this diff:
- tools = [_get_tool(tool_name) for tool_name in tool_names] - tools = await asyncio.gather(*tools) + tools = await asyncio.gather(*(_get_tool(tool_name) for tool_name in tool_names)) @@ - tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) + # Use already-fetched `tools` for ToolNodeAlso applies to: 107-107
docs/source/workflows/function-groups.md (3)
20-22: Use official naming on first mention.Replace “NeMo Agent toolkit” with “NVIDIA NeMo Agent toolkit” on first use, per guidelines.
-Function groups let you package multiple related functions together so they can share configuration, context, and resources within the NeMo Agent toolkit. +Function groups let you package multiple related functions together so they can share configuration, context, and resources within the NVIDIA NeMo Agent toolkit.
70-72: Fix incorrect import of Builder in example.Builder is defined in nat.builder.builder, not workflow_builder.
-from nat.builder.workflow_builder import Builder +from nat.builder.builder import Builder
173-189: Make filter callbacks async in docs.Filters are now async; update the example accordingly.
- def math_filter(function_names): - # Only allow functions that start with "add" - return [name for name in function_names if name.startswith("add")] + async def math_filter(function_names): + # Only allow functions that start with "add" + return [name for name in function_names if name.startswith("add")]tests/nat/builder/test_builder.py (1)
921-934: Avoid variable shadowing; fix tautological assertion.workflow_config variable is overwritten, making “workflow_config.workflow == workflow_config.workflow” a tautology.
- workflow = await builder.build() - - workflow_config = workflow.config - - assert workflow_config.general == general_config - assert workflow_config.functions == {"function1": function_config} - assert workflow_config.workflow == workflow_config.workflow + workflow = await builder.build() + + built_cfg = workflow.config + + assert built_cfg.general == general_config + assert built_cfg.functions == {"function1": function_config} + assert built_cfg.workflow == built_cfg.workflow # consider asserting equality to the set workflow config if availablesrc/nat/experimental/decorators/experimental_warning_decorator.py (1)
64-77: Overloads look good; consider ParamSpec for call signature preservation.Optional: use ParamSpec to preserve arg types/return types on wrapped callables for better type inference.
Example pattern:
from typing import ParamSpec P = ParamSpec("P") R = TypeVar("R") @overload def experimental(func: Callable[P, R], *, feature_name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[P, R]: ... @overload def experimental(*, feature_name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[Callable[P, R]], Callable[P, R]]: ...src/nat/builder/eval_builder.py (1)
94-108: Add return type annotation to get_all_tools.Per guidelines, include return typing.
- async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str): + async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[object]:If you prefer, import typing.Any and use list[Any].
src/nat/control_flow/sequential_executor.py (1)
72-75: Correct misleading debug log before compatibility check.The log states “is not compatible” before running the check. Reword to avoid false negatives in logs.
Apply:
- logger.debug( - f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with" - f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}") + logger.debug("Checking type compatibility: %s (out=%s) -> %s (in=%s)", + src_fn.instance_name, src_output_type, target_fn.instance_name, target_input_type)tests/nat/agent/test_reasoning_agent.py (1)
174-178: Silence ruff ARG001: prefix unused params with underscores.Several local async mocks don’t use their parameters. Prefix with
_to satisfy ruff.Apply this pattern here and similarly for other mocks in this file:
- async def mock_get_llm(name, wrapper_type): + async def mock_get_llm(_name, _wrapper_type): return mock_llmAlso update analogous mocks:
mock_get_function, etc. (see Lines 131, 163, 211, 220, 252, 270, 287, 306, 330, 353, 367, 389, 405).packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (2)
136-143: AcceptFunctionRefinget_function.Match
Builderand handle refs by stringifying.Apply:
- async def get_function(self, name: str) -> Function: + async def get_function(self, name: str | FunctionRef) -> Function: """Return a mock function if one is configured.""" + if isinstance(name, FunctionRef): + name = str(name) if name in self._mocks: mock_fn = AsyncMock() mock_fn.ainvoke = AsyncMock(return_value=self._mocks[name]) return mock_fn raise ValueError(f"Function '{name}' not mocked. Use mock_function() to add it.")
152-159: AcceptFunctionGroupRefinget_function_group.Match
Builderand handle refs by stringifying.Apply:
- async def get_function_group(self, name: str) -> FunctionGroup: + async def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup: """Return a mock function group if one is configured.""" + if isinstance(name, FunctionGroupRef): + name = str(name) if name in self._mocks: mock_fn_group = MagicMock(spec=FunctionGroup) mock_fn_group.ainvoke = AsyncMock(return_value=self._mocks[name]) return mock_fn_group raise ValueError(f"Function group '{name}' not mocked. Use mock_function_group() to add it.")src/nat/builder/workflow_builder.py (2)
545-547: Normalize refs before duplicate check inget_tools.
"foo"andFunctionRef("foo")should be considered the same logical tool; normalize to strings before computing uniqueness.Apply:
- unique = set(tool_names) - if len(unique) != len(tool_names): + normalized_names = [str(n) if isinstance(n, (FunctionRef, FunctionGroupRef)) else n for n in tool_names] + unique = set(normalized_names) + if len(unique) != len(normalized_names): raise ValueError("Tool names must be unique")
1176-1185: Fix dependency tracking for tool names (FunctionRef vs str).Membership test against
_function_groupsmust use string keys; also record stringified names.Apply:
- tools = await self._workflow_builder.get_tools(tool_names, wrapper_type) - for tool_name in tool_names: - if tool_name in self._workflow_builder._function_groups: - self._dependencies.add_function_group(tool_name) - else: - self._dependencies.add_function(tool_name) + tools = await self._workflow_builder.get_tools(tool_names, wrapper_type) + for tool_name in tool_names: + name_str = str(tool_name) + if name_str in self._workflow_builder._function_groups: + self._dependencies.add_function_group(name_str) + else: + self._dependencies.add_function(name_str)src/nat/profiler/decorators/function_tracking.py (5)
53-57: **Annotate *args and kwargs.Add explicit Any annotations to satisfy “type hints for all parameters” and improve editor/pyright help.
-def _prepare_serialized_args_kwargs(*args, **kwargs) -> tuple[list[Any], dict[str, Any]]: +def _prepare_serialized_args_kwargs(*args: Any, **kwargs: Any) -> tuple[list[Any], dict[str, Any]]:
60-68: Tighten push_intermediate_step parameter types.args/kwargs here are always serialized to list/dict; reflect that in the signature.
-def push_intermediate_step(step_manager: IntermediateStepManager, - identifier: str, - function_name: str, - event_type: IntermediateStepType, - args: Any = None, - kwargs: Any = None, - output: Any = None, - metadata: dict[str, Any] | None = None) -> None: +def push_intermediate_step(step_manager: IntermediateStepManager, + identifier: str, + function_name: str, + event_type: IntermediateStepType, + args: list[Any] | None = None, + kwargs: dict[str, Any] | None = None, + output: Any = None, + metadata: dict[str, Any] | None = None) -> None:
99-109: Use Callable in the implementation signature.Align the concrete signature with the overloads for better type checking while keeping return as Any to stay compatible with overloads.
-def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None) -> Any: +def track_function(func: Callable[..., Any] | None = None, *, metadata: dict[str, Any] | None = None) -> Any:
132-172: Consider error-path instrumentation.If the wrapped function/generator raises, SPAN_END isn’t emitted. If your event stream expects balanced START/END (and/or error records), consider try/except/finally to record an END or error detail before re‑raising.
Also applies to: 178-206, 211-246, 248-275
16-24: Add a short module docstring.Guidelines require Google‑style docstrings for public modules. Add a brief description after the license header.
src/nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/nat/builder/builder.py (2)
152-159: Add return type to get_auth_providers.Public API in src/ must have annotated returns. Implementation returns a list.
Apply:
- async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]): + async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]) -> list[AuthProviderBase]:
225-233: Align get_retrievers typing with get_retriever overloads.Return type varies with wrapper_type. Current annotation always lists Retriever, which is inaccurate when wrapper_type is provided.
Apply:
- async def get_retrievers(self, - retriever_names: Sequence[str | RetrieverRef], - wrapper_type: LLMFrameworkEnum | str | None = None) -> list[Retriever]: + async def get_retrievers(self, + retriever_names: Sequence[str | RetrieverRef], + wrapper_type: LLMFrameworkEnum | str | None = None) -> list[typing.Any]:Additionally, mirror overloads (add these above this method):
@typing.overload async def get_retrievers(self, retriever_names: Sequence[str | RetrieverRef], wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]: ... @typing.overload async def get_retrievers(self, retriever_names: Sequence[str | RetrieverRef], wrapper_type: None = None) -> list[Retriever]: ...src/nat/builder/workflow_builder.py (1)
1303-1307: Fix get_retriever return type; may return wrapper when wrapper_type provided.Signature promises Retriever but returns a framework wrapper when wrapper_type is set.
Apply:
- async def get_retriever(self, retriever_name: str, wrapper_type: LLMFrameworkEnum | str | None = None) -> Retriever: + async def get_retriever(self, retriever_name: str, wrapper_type: LLMFrameworkEnum | str | None = None) -> typing.Any:Optional: mirror overloads from the base interface for full type precision.
🧹 Nitpick comments (2)
src/nat/builder/workflow_builder.py (2)
1176-1185: Track dependencies using normalized names (FunctionRef vs str).Membership checks against _function_groups/_functions use string keys; current code may miss dependencies when refs are passed.
Apply:
- for tool_name in tool_names: - if tool_name in self._workflow_builder._function_groups: - self._dependencies.add_function_group(tool_name) - else: - self._dependencies.add_function(tool_name) + for tool_name in tool_names: + name_str = str(tool_name) if isinstance(tool_name, (FunctionRef, FunctionGroupRef)) else tool_name + if name_str in self._workflow_builder._function_groups: + self._dependencies.add_function_group(name_str) + else: + self._dependencies.add_function(name_str)
1188-1194: Normalize fn_name before recording dependency.Avoids missing deps when a FunctionRef is provided.
Apply:
- fn = await self._workflow_builder.get_tool(fn_name, wrapper_type) - - self._dependencies.add_function(fn_name) + fn = await self._workflow_builder.get_tool(fn_name, wrapper_type) + name_str = str(fn_name) if isinstance(fn_name, FunctionRef) else fn_name + self._dependencies.add_function(name_str)
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/nat/builder/builder.py(10 hunks)src/nat/builder/workflow_builder.py(23 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/builder/workflow_builder.pysrc/nat/builder/builder.py
src/nat/**/*
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Core functionality under src/nat should prioritize backward compatibility when changed
Files:
src/nat/builder/workflow_builder.pysrc/nat/builder/builder.py
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/builder/workflow_builder.pysrc/nat/builder/builder.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/builder/workflow_builder.pysrc/nat/builder/builder.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/builder/workflow_builder.pysrc/nat/builder/builder.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/builder/workflow_builder.pysrc/nat/builder/builder.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/builder/workflow_builder.pysrc/nat/builder/builder.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/builder/workflow_builder.pysrc/nat/builder/builder.py
🧠 Learnings (1)
📚 Learning: 2025-09-09T20:32:39.016Z
Learnt from: CR
PR: NVIDIA/NeMo-Agent-Toolkit#0
File: .cursor/rules/nat-test-llm.mdc:0-0
Timestamp: 2025-09-09T20:32:39.016Z
Learning: Applies to **/*.py : When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.<FRAMEWORK>) and call the framework’s method (e.g., ainvoke, achat, call).
Applied to files:
src/nat/builder/workflow_builder.pysrc/nat/builder/builder.py
🧬 Code graph analysis (2)
src/nat/builder/workflow_builder.py (3)
src/nat/builder/workflow.py (1)
Workflow(41-160)src/nat/builder/function.py (5)
get_included_functions(580-626)Function(46-301)FunctionGroup(351-713)get_accessible_functions(485-521)add_function(380-428)src/nat/builder/builder.py (18)
get_function(76-77)get_function_group(80-81)get_tools(112-115)get_tool(118-119)add_llm(122-123)get_llm(126-127)add_retriever(222-223)add_ttc_strategy(260-261)add_memory_client(200-201)add_object_store(161-162)add_function_group(72-73)add_function(68-69)add_auth_provider(144-146)get_auth_provider(149-150)get_retriever(236-238)get_retriever(241-242)get_retriever(245-246)get_retriever(249-252)
src/nat/builder/builder.py (3)
src/nat/experimental/decorators/experimental_warning_decorator.py (3)
experimental(66-68)experimental(72-74)experimental(77-151)src/nat/builder/workflow_builder.py (24)
get_function(466-472)get_function(1138-1144)get_function_group(475-481)get_function_group(1147-1153)get_tools(541-572)get_tools(1176-1185)get_tool(575-589)get_tool(1188-1194)add_llm(592-605)add_llm(1197-1198)add_auth_provider(639-677)add_auth_provider(1202-1203)add_embedder(708-721)add_embedder(1222-1223)add_memory_client(754-765)add_memory_client(1238-1239)get_memory_client(768-775)get_memory_client(1242-1250)add_retriever(814-828)add_retriever(1299-1300)add_ttc_strategy(864-877)add_ttc_strategy(1277-1278)get_ttc_strategy(880-906)get_ttc_strategy(1281-1287)src/nat/builder/eval_builder.py (2)
get_tool(97-103)get_all_tools(94-107)
🪛 Ruff (0.13.1)
src/nat/builder/workflow_builder.py
457-457: Avoid specifying long messages outside the exception class
(TRY003)
547-547: Avoid specifying long messages outside the exception class
(TRY003)
557-557: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (6)
src/nat/builder/builder.py (4)
83-89: LGTM: Async batch getters with asyncio.gather are correct.Order-preserving, exceptions propagate, and return types align.
142-146: Decorator order for abstract experimental API looks correct.@AbstractMethod is outermost, preserving abstractness while keeping experimental warnings.
258-261: Decorator order is correct here as well.Abstractness preserved; experimental warning applied.
121-123: Fix return type of add_llm to None (matches implementations).Concrete builders return None. Keeping typing.Any here breaks type checking and misleads users.
Apply:
- async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> typing.Any: + async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> None:src/nat/builder/workflow_builder.py (2)
220-335: LGTM: build is now async and correctly awaits async function-group accessors.Entry-function resolution and component wiring are consistent; no blocking I/O on sync paths.
454-461: Correctly awaiting FunctionGroup.get_included_functions and merging into registry.Conflict detection with existing functions is sound.
willkill07
force-pushed
the
wkk_async-function-group-callback
branch
from
c8ae89f to
8acce7b
Compare
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
willkill07
force-pushed
the
wkk_async-function-group-callback
branch
from
8acce7b to
9c994fa
Compare
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (18)
src/nat/agent/tool_calling_agent/register.py (2)
52-52: Add return type and function docstring (public API; pyright treats missing return types as errors).This is an async generator that yields
FunctionInfo. Add-> AsyncIterator[FunctionInfo]and a short Google‑style docstring.Apply:
-async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, builder: Builder): +async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, builder: Builder) -> AsyncIterator[FunctionInfo]: + """Register the tool-calling agent workflow. + + Yields: + `FunctionInfo`: Async callable entrypoint that runs the Tool Calling Agent. + """And add import:
+from collections.abc import AsyncIterator
41-41: Fix spelling in field description.“stoping” → “stopping”.
Apply:
- max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.") + max_iterations: int = Field(default=15, description="Number of tool calls before stopping the tool calling agent.")packages/nvidia_nat_test/src/nat/test/tool_test_runner.py (3)
136-142: Fix parameter type and return type to match Builder interface.The method signature should match the parent
Builderclass which acceptsstr | FunctionRefand returnsFunction.Apply this diff to align with the parent interface:
- async def get_function(self, name: str) -> Function: + async def get_function(self, name: str | FunctionRef) -> Function: """Return a mock function if one is configured.""" + if isinstance(name, FunctionRef): + name = str(name) if name in self._mocks: mock_fn = AsyncMock() mock_fn.ainvoke = AsyncMock(return_value=self._mocks[name]) return mock_fn raise ValueError(f"Function '{name}' not mocked. Use mock_function() to add it.")Add the required import:
from nat.data_models.component_ref import FunctionRef
152-158: Fix parameter type to match Builder interface.The method signature should accept
str | FunctionGroupRefto match the parentBuilderclass.Apply this diff:
- async def get_function_group(self, name: str) -> FunctionGroup: + async def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup: """Return a mock function group if one is configured.""" + if isinstance(name, FunctionGroupRef): + name = str(name) if name in self._mocks: mock_fn_group = MagicMock(spec=FunctionGroup) mock_fn_group.ainvoke = AsyncMock(return_value=self._mocks[name]) return mock_fn_group raise ValueError(f"Function group '{name}' not mocked. Use mock_function_group() to add it.")Add the required import:
from nat.data_models.component_ref import FunctionGroupRef
227-235: Fix parameter type to match Builder interface.The method should accept
str | MemoryRefto align with the parentBuilderclass.Apply this diff:
- async def get_memory_client(self, memory_name: str) -> MemoryEditor: + async def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor: """Return a mock memory client if one is configured.""" + if isinstance(memory_name, MemoryRef): + memory_name = str(memory_name) key = f"memory_{memory_name}" if key in self._mocks: mock_memory = MagicMock() mock_memory.add = AsyncMock(return_value=self._mocks[key]) mock_memory.search = AsyncMock(return_value=self._mocks[key]) return mock_memory raise ValueError(f"Memory client '{memory_name}' not mocked. Use mock_memory_client() to add it.")docs/source/extend/functions.md (2)
71-73: Fix first-use branding: “NVIDIA NeMo Agent toolkit”.Per docs guidelines, the first mention must say “NVIDIA NeMo Agent toolkit”.
Apply this diff:
- To use a function from a configuration file, it must be registered with NeMo Agent toolkit. Registering a function is done with the {py:deco}`nat.cli.register_workflow.register_function` decorator. More information about registering components can be found in the [Plugin System](../extend/plugins.md) documentation. + To use a function from a configuration file, it must be registered with NVIDIA NeMo Agent toolkit. Registering a function is done with the {py:deco}`nat.cli.register_workflow.register_function` decorator. More information about registering components can be found in the [Plugin System](../extend/plugins.md) documentation.
218-219: Docs bug: undefined variablemy_callable.The snippet defines
_my_functionbut yieldsmy_callable.Apply this diff:
- yield my_callable + yield _my_functionexamples/frameworks/multi_frameworks/src/nat_multi_frameworks/register.py (1)
116-123: Router returns unmapped value; fix typo to avoid runtime error.Router returns
"supevisor"which isn’t in the conditional edge mapping; this will fail if hit.Apply this diff to route safely to a mapped edge:
- route_to = "supevisor" + route_to = "workers"src/nat/experimental/decorators/experimental_warning_decorator.py (1)
92-99: Docstring nit: correct decorator name.Mentions
@track_function(...); should be@experimental(...).Apply this diff:
- # If called as @track_function(...) but not immediately passed a function + # If called as @experimental(...) but not immediately passed a functionsrc/nat/profiler/parameter_optimization/prompt_optimizer.py (3)
229-245: Semaphore doesn’t bound the heavy evaluation; moveasync with sem:to coverrun_and_evaluate().Currently only
apply_suggestionsis inside the semaphore, so concurrent evaluations can overwhelm the endpoint. Wrap the whole evaluation (including reps) under the semaphore to honorga_parallel_evaluations.async def _evaluate(ind: Individual) -> Individual: - async with sem: - cfg_trial = apply_suggestions(base_cfg, ind.prompts) - eval_cfg = EvaluationRunConfig( - config_file=cfg_trial, - dataset=opt_run_config.dataset, - result_json_path=opt_run_config.result_json_path, - endpoint=opt_run_config.endpoint, - endpoint_timeout=opt_run_config.endpoint_timeout, - override=opt_run_config.override, - ) - # Run reps sequentially under the same semaphore to avoid overload - all_results: list[list[tuple[str, Any]]] = [] - for _ in range(reps): - res = (await EvaluationRun(config=eval_cfg).run_and_evaluate()).evaluation_results - all_results.append(res) + # Bound end-to-end evaluation concurrency to avoid overload + async with sem: + cfg_trial = apply_suggestions(base_cfg, ind.prompts) + eval_cfg = EvaluationRunConfig( + config_file=cfg_trial, + dataset=opt_run_config.dataset, + result_json_path=opt_run_config.result_json_path, + endpoint=opt_run_config.endpoint, + endpoint_timeout=opt_run_config.endpoint_timeout, + override=opt_run_config.override, + ) + # Run reps sequentially under the same semaphore to avoid overload + all_results: list[list[tuple[str, Any]]] = [] + for _ in range(reps): + res = (await EvaluationRun(config=eval_cfg).run_and_evaluate()).evaluation_results + all_results.append(res)
344-355: Honorga_offspring_size; remove no-op loop and bound offspring count.
for _ in range(max(0, offspring_size), needed): passis a no-op, andoffspring_sizeis ignored. Respect it when producing children.- offspring: list[Individual] = [] - for _ in range(max(0, offspring_size), needed): - pass # ensure bound correctness - while len(offspring) < needed: + offspring: list[Individual] = [] + # Cap children by configured offspring_size while not exceeding population needs + num_children = max(0, min(int(offspring_size), needed)) + while len(offspring) < num_children: p1 = _select_parent(population) p2 = _select_parent(population) if p2 is p1 and len(population) > 1: p2 = random.choice([ind for ind in population if ind is not p1]) child = await _make_child(p1, p2) offspring.append(child) population = next_population + offspring
212-215: Uselogger.exception()when catching without re-raising to keep stack traces.Per guidelines: when not re-raising, log with
logger.exception().- except Exception as e: - logger.warning("Mutation failed for %s: %s; using original.", param, e) + except Exception: + logger.exception("Mutation failed for %s; using original.", param) ... - except Exception as e: - logger.warning("Recombination failed for %s: %s; falling back to parent.", param, e) + except Exception: + logger.exception("Recombination failed for %s; falling back to parent.", param) child = random.choice([pa, pb]) ... - except Exception as e: - logger.warning("Mutation failed for %s: %s; keeping child as-is.", param, e) + except Exception: + logger.exception("Mutation failed for %s; keeping child as-is.", param) ... - except Exception as e: # pragma: no cover - best effort - logger.warning("Failed to write GA history CSV: %s", e) + except Exception: # pragma: no cover - best effort + logger.exception("Failed to write GA history CSV")Also applies to: 283-294, 379-381
tests/nat/front_ends/mcp/test_mcp_custom_routes.py (1)
61-66: Mark async test with@pytest.mark.asyncio.Without the marker (unless auto mode is configured), the async test may not run.
-async def test_custom_mcp_worker(mcp_nat_config: Config): +@pytest.mark.asyncio +async def test_custom_mcp_worker(mcp_nat_config: Config):src/nat/control_flow/sequential_executor.py (1)
72-80: Fix misleading debug log; emit detail only when incompatible.The current debug message states “is not compatible” before the check.
- logger.debug( - f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with" - f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}") + logger.debug("Checking type compatibility: %s (%s) -> %s (%s)", + src_fn.instance_name, str(src_output_type), + target_fn.instance_name, str(target_input_type)) is_compatible = DecomposedType.is_type_compatible(src_output_type, target_input_type) if not is_compatible: raise ValueError( f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with" f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}")docs/source/workflows/function-groups.md (2)
20-21: Use the full toolkit name on first mention.Per docs guidelines, first use should be “NVIDIA NeMo Agent toolkit”; subsequent mentions can use “NeMo Agent toolkit.”
Apply:
-Function groups let you package multiple related functions together so they can share configuration, context, and resources within the NeMo Agent toolkit. +Function groups let you package multiple related functions together so they can share configuration, context, and resources within the NVIDIA NeMo Agent toolkit.
173-186: Make the filter callback async in the example.Function-group filter callbacks are async now. Update
math_filtertoasync defto match the API.Apply:
- def math_filter(function_names): - # Only allow functions that start with "add" - return [name for name in function_names if name.startswith("add")] + async def math_filter(function_names): + # Only allow functions that start with "add" + return [name for name in function_names if name.startswith("add")]src/nat/builder/workflow_builder.py (1)
1303-1307: Fix return annotation for ChildBuilder.get_retriever.When
wrapper_typeis provided, return isn’tRetriever. Mirror overloads or usetyping.Any.Apply one of:
- Simplest:
- async def get_retriever(self, retriever_name: str, wrapper_type: LLMFrameworkEnum | str | None = None) -> Retriever: + async def get_retriever(self, retriever_name: str, wrapper_type: LLMFrameworkEnum | str | None = None) -> typing.Any:
- Or add overloads identical to interface for precise typing.
src/nat/profiler/decorators/function_tracking.py (1)
132-172: Always close spans on exceptions (avoid dangling spans).
Exceptions in the current wrappers skip SPAN_END — wrap the call/iteration in try/except to push SPAN_END with a serialized exception payload (or emit a dedicated IntermediateStepType.ERROR then SPAN_END if that enum value exists), then re-raise.
Applies: src/nat/profiler/decorators/function_tracking.py — async gen (132–172), async fn (174–205), sync gen (211–246), sync fn (248–275).
🧹 Nitpick comments (17)
src/nat/agent/tool_calling_agent/register.py (2)
16-16: Add a module-level docstring.Public modules must have a concise Google‑style docstring.
Apply:
+"""Tool-calling agent workflow registration and runtime entrypoint."""
34-36: Minor capitalization consistency in class docstring.Use “Tool Calling Agent” consistently.
Apply:
- A Tool Calling Agent requires an LLM which supports tool calling. A tool Calling Agent utilizes the tool + A Tool Calling Agent requires an LLM which supports tool calling. A Tool Calling Agent utilizes the tooldocs/source/extend/functions.md (1)
197-205: Grammar nit: “It’s required …”.Minor copy fix.
Apply this diff:
-Its required to use an async context manager coroutine to register a function +It's required to use an async context manager coroutine to register a function.cursor/rules/nat-workflows/add-functions.mdc (1)
1-5: Frontmatter: remove emptyglobs:or provide a list.
globs:is present without values; either drop it or set concrete globs.Apply this diff to remove the unused key:
---- -description: Follow these rules when the user's request involves implementing, adding, creating, or modifying functions within NeMo Agent Toolkit workflows -globs: -alwaysApply: false ---- +--- +description: Follow these rules when the user's request involves implementing, adding, creating, or modifying functions within NeMo Agent Toolkit workflows +alwaysApply: false +---examples/frameworks/multi_frameworks/src/nat_multi_frameworks/register.py (1)
66-73: Typo in prompt string (“Classifcation”).Minor user-facing text fix.
Apply this diff:
- Classifcation topic: + Classification topic:examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py (1)
86-95: Deduplicate tool retrieval; useget_toolsonce and reuse.You fetch tools via
gatherand then again viabuilder.get_tools. Consolidate to one call.Apply this diff:
- async def _get_tool(tool_name: str): - return await builder.get_tool(tool_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN) - - tools = [_get_tool(tool_name) for tool_name in tool_names] - tools = await asyncio.gather(*tools) + tools = await builder.get_tools(tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN) llm_n_tools = llm.bind_tools(tools, parallel_tool_calls=True) - categorizer_tool = await _get_tool("categorizer") - maintenance_check_tool = await _get_tool("maintenance_check") + categorizer_tool = await builder.get_tool("categorizer", wrapper_type=LLMFrameworkEnum.LANGCHAIN) + maintenance_check_tool = await builder.get_tool("maintenance_check", wrapper_type=LLMFrameworkEnum.LANGCHAIN) @@ - tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)Also applies to: 106-108
src/nat/experimental/decorators/experimental_warning_decorator.py (1)
31-31: Type annotate the issued-warnings set.Helps static analysis.
Apply this diff:
-_warning_issued = set() +_warning_issued: set[str] = set()src/nat/profiler/parameter_optimization/prompt_optimizer.py (1)
45-52: Add a concise docstring foroptimize_prompts(public API).Public functions must have Google‑style docstrings.
async def optimize_prompts( *, base_cfg: Config, full_space: dict[str, SearchSpace], optimizer_config: OptimizerConfig, opt_run_config: OptimizerRunConfig, ) -> None: + """ + Run GA-based prompt optimization over the provided prompt search space. + + Args: + base_cfg: Base NAT configuration used to build/evaluate workflows. + full_space: Full parameter search space; prompts are selected via `is_prompt`. + optimizer_config: GA and evaluation metric configuration. + opt_run_config: Dataset, endpoint, and output paths for evaluation/results. + """src/nat/experimental/test_time_compute/functions/plan_select_execute_function.py (1)
107-110: Fetch tool metadata concurrently for performance.Multiple awaited lookups in a loop serialize I/O. Use
asyncio.gather.- for tool in function_used_tools: - tool_impl = await builder.get_function(tool) - tool_list += f"- {tool}: {tool_impl.description if hasattr(tool_impl, 'description') else ''}\n" + # Resolve tool implementations concurrently + import asyncio # safe local import if not yet imported at module top + tool_impls = await asyncio.gather(*[builder.get_function(tool) for tool in function_used_tools]) + for tool, tool_impl in zip(function_used_tools, tool_impls): + tool_list += f"- {tool}: {getattr(tool_impl, 'description', '')}\n"tests/nat/builder/test_function_group.py (1)
149-154: Silence Ruff ARG001 by prefixing unused parameters with underscores.Keeps signatures while satisfying linter.
- async def exclude_func1(name: str) -> bool: + async def exclude_func1(_name: str) -> bool: return False # Always exclude func1 - - async def include_func2(name: str) -> bool: + + async def include_func2(_name: str) -> bool: return True # Always include func2 ... - async def exclude_func2(name: str) -> bool: + async def exclude_func2(_name: str) -> bool: return False ... - async def exclude_func1(name: str) -> bool: + async def exclude_func1(_name: str) -> bool: return False ... - async def empty_filter(names: Sequence[str]) -> Sequence[str]: + async def empty_filter(_names: Sequence[str]) -> Sequence[str]: return [] ... - async def exclude_filter(name: str) -> bool: + async def exclude_filter(_name: str) -> bool: return False ... - async def include_filter(name: str) -> bool: + async def include_filter(_name: str) -> bool: return TrueAlso applies to: 152-154, 236-238, 302-304, 363-365, 616-618, 623-625
src/nat/agent/reasoning_agent/reasoning_agent.py (1)
121-124: Resolve tool implementations concurrently to avoid N+1 latency.- for tool in function_used_tools: - tool_impl = await builder.get_function(tool) - tool_names_with_desc.append((tool, tool_impl.description if hasattr(tool_impl, "description") else "")) + import asyncio # add at module top if preferred + impls = await asyncio.gather(*[builder.get_function(tool) for tool in function_used_tools]) + for tool, tool_impl in zip(function_used_tools, impls): + tool_names_with_desc.append((tool, getattr(tool_impl, "description", "")))src/nat/control_flow/sequential_executor.py (1)
125-127: Preserve stack traces on unexpected errors.Avoid wrapping unknown exceptions in a new
ValueErrorwithout chaining.- except Exception as e: - raise ValueError(f"Error with the sequential executor tool list: {e}") + except Exception: + # Keep original type/trace + raisesrc/nat/builder/eval_builder.py (1)
16-16: Import typing for explicit return annotations added below.Apply:
-import asyncio +import asyncio +import typingsrc/nat/builder/function.py (1)
460-469: Optional: deduplicate identity filter creation.Extract a small
_identity_filterasync helper to avoid repeating definitions.Apply:
+ @staticmethod + async def _identity_filter(x: Sequence[str]) -> Sequence[str]: + return xThen replace local identity_filter definitions with
self._identity_filter.Also applies to: 552-561, 611-620, 649-658
src/nat/profiler/decorators/function_tracking.py (3)
86-97: Prefer ParamSpec + ReturnType for precise callable signatures.This improves compatibility with pyright/mypy and preserves parameter types and return type more explicitly than a Callable‑bound
TypeVar.Apply:
-from typing import TypeVar +from typing import TypeVar +from typing import ParamSpec @@ -F = TypeVar('F', bound=Callable[..., Any]) +F = TypeVar('F', bound=Callable[..., Any]) +P = ParamSpec('P') +R = TypeVar('R') @@ -@overload -def track_function(func: F, *, metadata: dict[str, Any] | None = None) -> F: +@overload +def track_function(func: Callable[P, R], *, metadata: dict[str, Any] | None = None) -> Callable[P, R]: @@ -@overload -def track_function(*, metadata: dict[str, Any] | None = None) -> Callable[[F], F]: +@overload +def track_function(*, metadata: dict[str, Any] | None = None) -> Callable[[Callable[P, R]], Callable[P, R]]:
279-289: Mirror typing precision fortrack_unregistered_function.Same ParamSpec/ReturnType approach recommended here to preserve callable signatures.
Apply:
@@ -@overload -def track_unregistered_function(func: F, *, name: str | None = None, metadata: dict[str, Any] | None = None) -> F: +@overload +def track_unregistered_function(func: Callable[P, R], *, name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[P, R]: @@ -@overload -def track_unregistered_function(*, name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[F], F]: +@overload +def track_unregistered_function(*, name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[Callable[P, R]], Callable[P, R]]:
99-108: Docstring completeness: add Args/Returns/Raises sections (Google style).Public APIs should include structured sections to aid users and pass docstyle checks.
Example for
track_function:
- Args: func, metadata
- Returns: Wrapped callable preserving type of
func- Raises: TypeError (invalid metadata)
Also applies to: 295-307
Signed-off-by: Will Killian <wkillian@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
src/nat/builder/builder.py (2)
86-93: Validate uniqueness (match get_tools) to avoid duplicate work/ambiguous results.Align with get_tools’ behavior by rejecting duplicate names.
async def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]: - tasks = [self.get_function(name) for name in function_names] + names = list(function_names) + if len(set(names)) != len(names): + raise ValueError("Function names must be unique") + tasks = [self.get_function(name) for name in names] return list(await asyncio.gather(*tasks, return_exceptions=False)) async def get_function_groups(self, function_group_names: Sequence[str | FunctionGroupRef]) -> list[FunctionGroup]: - tasks = [self.get_function_group(name) for name in function_group_names] + names = list(function_group_names) + if len(set(names)) != len(names): + raise ValueError("Function group names must be unique") + tasks = [self.get_function_group(name) for name in names] return list(await asyncio.gather(*tasks, return_exceptions=False))Also applies to: 90-93
79-85: Add concise Google-style docstrings to public APIs.Per coding guidelines, add docstrings to these public methods (examples above). Sample:
async def get_function(self, name: str | FunctionRef) -> Function: """Return the `Function` instance for `name` or raise if not found."""Also applies to: 86-93, 115-118, 121-127, 129-131, 132-140, 142-144, 146-149, 164-166, 182-184, 185-193, 195-201, 203-205, 225-227, 228-237, 252-255, 258-260, 262-264, 267-271, 274-279, 316-317
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/nat/builder/builder.py(10 hunks)
🧰 Additional context used
📓 Path-based instructions (6)
**/*.{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/builder/builder.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.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).
**/*.py: In code comments/identifiers use NAT abbreviations as specified: nat for API namespace/CLI, nvidia-nat for package name, NAT for env var prefixes; do not use these abbreviations in documentation
Follow PEP 20 and PEP 8; run yapf with column_limit=120; use 4-space indentation; end files with a single trailing newline
Run ruff check --fix as linter (not formatter) using pyproject.toml config; fix warnings unless explicitly ignored
Respect naming: snake_case for functions/variables, PascalCase for classes, UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: use bare raise to re-raise; log with logger.error() when re-raising to avoid duplicate stack traces; use logger.exception() when catching without re-raising
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise and ending with a period; surround code entities with backticks
Validate and sanitize all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile or mprof before optimizing; cache expensive computations with functools.lru_cache or external cache; leverage NumPy vectorized operations when beneficial
Files:
src/nat/builder/builder.py
src/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All importable Python code must live under src/ (or packages//src/)
Files:
src/nat/builder/builder.py
src/nat/**/*
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Changes in src/nat should prioritize backward compatibility
Files:
src/nat/builder/builder.py
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/builder/builder.py
{src/**/*.py,packages/*/src/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All public APIs must have Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Files:
src/nat/builder/builder.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/builder/builder.py
🧠 Learnings (1)
📚 Learning: 2025-09-09T20:32:39.016Z
Learnt from: CR
PR: NVIDIA/NeMo-Agent-Toolkit#0
File: .cursor/rules/nat-test-llm.mdc:0-0
Timestamp: 2025-09-09T20:32:39.016Z
Learning: Applies to **/*.py : When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.<FRAMEWORK>) and call the framework’s method (e.g., ainvoke, achat, call).
Applied to files:
src/nat/builder/builder.py
🧬 Code graph analysis (1)
src/nat/builder/builder.py (4)
src/nat/experimental/decorators/experimental_warning_decorator.py (3)
experimental(66-68)experimental(72-74)experimental(77-151)src/nat/experimental/test_time_compute/models/strategy_base.py (2)
StrategyBase(26-67)stage_type(58-60)src/nat/builder/workflow_builder.py (26)
get_function(466-472)get_function(1138-1144)get_function_group(475-481)get_function_group(1147-1153)get_tools(541-572)get_tools(1176-1185)get_tool(575-589)get_tool(1188-1194)add_llm(592-605)add_llm(1197-1198)add_auth_provider(639-677)add_auth_provider(1202-1203)add_object_store(787-797)add_object_store(1257-1258)add_embedder(708-721)add_embedder(1222-1223)add_memory_client(754-765)add_memory_client(1238-1239)get_memory_client(768-775)get_memory_client(1242-1250)add_retriever(814-828)add_retriever(1299-1300)add_ttc_strategy(864-877)add_ttc_strategy(1277-1278)get_ttc_strategy(880-906)get_ttc_strategy(1281-1287)src/nat/builder/eval_builder.py (2)
get_tool(97-103)get_all_tools(94-107)
🔇 Additional comments (11)
src/nat/builder/builder.py (11)
55-57: LGTM: forward-ref import is correct.Using TYPE_CHECKING + string annotation keeps runtime deps light and satisfies type checkers.
115-123: Async tool API surface looks good.
146-149: Decorator order fixed — preserves abstractness.
164-166: Good: explicit ObjectStore return type for add_object_store.
182-184: Good: add_embedder returns None (consistent).
206-212: LGTM: concurrent fetch with asyncio.gather.
214-218: Nice: docstring added here; please replicate for other public methods.
228-237: LGTM: typed list[Retriever] and optional wrapper_type.
270-271: LGTM: narrowed return type to StrategyBase (via forward ref).
316-317: EvalBuilder async get_all_tools signature looks consistent.
125-126: Return type should be None (matches implementations).Abstract add_llm returns typing.Any but concrete impls return None; fix signature.
- async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> typing.Any: + async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> None:
|
/merge |
…oup filters `async` (NVIDIA#834) This PR aligns the expected usage of `get_*` calls from the builder to all be `async`. This was needed currently for `get_tools`, but it makes sense to logically extend to other components of the workflow. ### Breaking Changes * All `get_*` functions of FunctionGroup are now `async` * All `Builder.get_*` functions are now `async` (this is for current need and future-proofing) * `Builder.build` is now `async` ### Orthogonal Changes * filter callbacks for function groups are now async * type hinting for `@experimental` and `@track_function` This change enables filter callbacks for function groups being marked async. While enabling this feature, it was discovered that it propagates async all the way up to `Builder.build()` (resulting in a breaking change). This is worthwhile due to the asynchronous nature of filtering functions (e.g. Authentication, Dynamic Checks via endpoint queries). Since we were already making one breaking change, we might as well make more to reduce future breaking changes :) Closes ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit - Refactor - Major public API migration to async: build, function/group, tool/memory/retriever fetches and related accessors now return awaitables; batch retrievals made concurrent; some type annotations and overloads updated. - Documentation - Examples and guides updated to use await in code samples. - Examples - Agents and workflows updated to await tool/function retrievals; selective lazy initialization added; a config default adjusted. - Tests - Test suites converted to asyncio with AsyncMock/await usage to match the async surface. Authors: - Will Killian (https://github.com/willkill07) Approvers: - Zhongxuan (Daniel) Wang (https://github.com/zhongxuanwang-nv) - David Gardner (https://github.com/dagardner-nv) URL: NVIDIA#834 Signed-off-by: Yuchen Zhang <yuchenz@nvidia.com>
In a prior PR (#834), we unified the `async` interface across most builder patterns. We missed await-ing `get_memory_client` in the built-in memory_tools. Closes nvbugs-5563968 ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit - New Features - None - Bug Fixes - Fixed inconsistent default memory selection to prevent unexpected errors. - Resolved occasional failures when connecting to the memory service by ensuring proper asynchronous handling. - Refactor - Standardized memory tool configuration for add, delete, and get operations to improve consistency and reliability. - Documentation - Clarified descriptions for memory tools to improve user-facing text. Authors: - Will Killian (https://github.com/willkill07) Approvers: - Yuchen Zhang (https://github.com/yczhang-nv) URL: #935
3 participants
Description
This PR aligns the expected usage of
get_*calls from the builder to all beasync. This was needed currently forget_tools, but it makes sense to logically extend to other components of the workflow.Breaking Changes
get_*functions of FunctionGroup are nowasyncBuilder.get_*functions are nowasync(this is for current need and future-proofing)Builder.buildis nowasyncOrthogonal Changes
@experimentaland@track_functionThis change enables filter callbacks for function groups being marked async. While enabling this feature, it was discovered that it propagates async all the way up to
Builder.build()(resulting in a breaking change). This is worthwhile due to the asynchronous nature of filtering functions (e.g. Authentication, Dynamic Checks via endpoint queries).Since we were already making one breaking change, we might as well make more to reduce future breaking changes :)
Closes
By Submitting this PR I confirm:
Summary by CodeRabbit