Add AG2 (formerly AutoGen) plugin package and demo#1810
Open
VasiliyRad wants to merge 37 commits intoNVIDIA:developfrom
Open
Add AG2 (formerly AutoGen) plugin package and demo#1810VasiliyRad wants to merge 37 commits intoNVIDIA:developfrom
VasiliyRad wants to merge 37 commits intoNVIDIA:developfrom
Conversation
Merge develop to main for rc-10
Merge develop to main
Merge 1.0.0: upstream/develop to upstream/main
Merge pull request NVIDIA#92 from NVIDIA/develop
Updated changelog with another bug fix (NVIDIA#93)
[RELEASE] AIQToolkit v1.1.0
Changelog updated for 1.1.0 (NVIDIA#291)
## Description <!-- Note: The pull request title will be included in the CHANGELOG. --> <!-- Provide a standalone description of changes in this PR. --> <!-- Reference any issues closed by this PR with "closes NVIDIA#1234". All PRs should have an issue they close--> <!-- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> ## ❄️ Code freeze for `develop` and `v1.2.0` release ### What does this mean? Only critical/hotfix level issues should be merged into `develop` until release (merging of this PR). ### What is the purpose of this PR? - Update documentation - Allow testing for the new release - Enable a means to merge `develop` into `main` for the release ## 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.
<!-- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> ## ❄️ Code freeze for `v1.2.1` release ### What does this mean? Only critical/hotfix level issues should be merged into `develop` until release (merging of this PR). ### What is the purpose of this PR? - Update documentation - Allow testing for the new release - Enable a means to merge `develop` into `main` for the release
Signed-off-by: David Gardner <dagardner@nvidia.com>
<!-- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> ## ❄️ Code freeze for `release/1.3` and `v1.3.0` release ### What does this mean? Only critical/hotfix level issues should be merged into `release/1.3` until release (merging of this PR). ### What is the purpose of this PR? - Update documentation - Allow testing for the new release - Enable a means to merge `release/1.3` into `main` for the release
Signed-off-by: David Gardner <dagardner@nvidia.com>
<!-- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> ## ❄️ Code freeze for `release/1.3` and `v1.3.1` release ### What does this mean? Only critical/hotfix level issues should be merged into `release/1.3` until release (merging of this PR). ### What is the purpose of this PR? - Update documentation - Allow testing for the new release - Enable a means to merge `release/1.3` into `main` for the release
Signed-off-by: David Gardner <dagardner@nvidia.com>
<!-- SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> ## ❄️ Code freeze for `v1.4.0` release ### What does this mean? Only critical/hotfix level issues should be merged into `release/1.4` until release (merging of this PR). ### What is the purpose of this PR? - [RELEASE] NeMo-Agent-Toolkit v1.4.0
Signed-off-by: Michael Demoret <mdemoret@nvidia.com>
<!-- SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. --> ## ❄️ Code freeze for `develop` and `v1.5.0` release ### What does this mean? Only critical/hotfix level issues should be merged into `develop` until release (merging of this PR). ### What is the purpose of this PR? - Update documentation - Allow testing for the new release - Enable a means to merge `develop` into `main` for the release
Signed-off-by: Anuradha Karuppiah <26330987+AnuradhaKaruppiah@users.noreply.github.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds an AG2 (formerly AutoGen) integration: new nvidia-nat-ag2 plugin with LLM wrappers, tool wrapper, and profiler; example AG2 demo project (workflows, tools, configs, tests, README); packaging/entry-point additions; docs updated; and an Changes
Sequence Diagram(s)sequenceDiagram
participant Builder as Builder
participant NAT as NAT Core
participant AG2 as AG2 Plugin
participant Profiler as AG2 Profiler
participant LLM as LLM Provider
participant Tool as AG2 Tool/Function
Builder->>AG2: register LLM/tool wrappers & functions
Builder->>AG2: invoke ag2_team / ag2_async_team / ag2_research_team
AG2->>Tool: convert NAT Function -> AG2 Tool (ag2_tool_wrapper)
AG2->>NAT: schedule Function execution via Tool
Note over Profiler,LLM: Profiler instruments autogen.oai.client.OpenAIWrapper.create
Tool->>Profiler: call OpenAIWrapper.create(...) (patched)
Profiler->>LLM: emit LLM_START and call original create
LLM-->>Profiler: return output
Profiler->>Builder: emit LLM_END with timing/metadata
Tool->>AG2: return Function result
AG2-->>Builder: workflow returns final content
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (5)
docs/source/components/integrations/frameworks.md (1)
93-94: Minor: Consider varying sentence structure.Three successive sentences begin with "AG2." Consider rewording for better readability, e.g., "The framework provides a modern API for building collaborative agent systems..."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/source/components/integrations/frameworks.md` around lines 93 - 94, The paragraph contains three successive sentences that all start with "AG2", hurting readability; rewrite one or two of those sentences to vary sentence openings (for example start one with "The framework" or "It provides") and rephrase to keep the same meaning about community maintenance, multi-agent orchestration, and features (GroupChat patterns, tool integration, flexible LLM configuration and MCP integration) while ensuring the term "AG2" still appears at least once for clarity.packages/nvidia_nat_ag2/pyproject.toml (1)
14-18: Consider addingnvidia-nat-testas an optional test dependency.Per coding guidelines, packages should likely list
nvidia-nat-testas an optional dependency in thetestextra to support testing infrastructure.💡 Suggested addition
[tool.setuptools_dynamic_dependencies] dependencies = [ "nvidia-nat-core == {version}", "ag2[openai] >= 0.11.0", ] + +[project.optional-dependencies] +test = [ + "nvidia-nat-test == {version}", +]As per coding guidelines: "
nvidia-nat-testshould likely be listed as an optional dependency in thetestextra."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/pyproject.toml` around lines 14 - 18, Add `nvidia-nat-test` as an optional test dependency by defining a `test` extra and including the package there (e.g., add a `project.optional-dependencies` / `tool.setuptools_dynamic_dependencies` entry for the `test` extra containing "nvidia-nat-test >= {version}" or the appropriate version constraint); update the existing dependencies block (referenced by the `dependencies` symbol) to keep runtime deps unchanged and add the new `test` extra so CI and local test tooling can install `nvidia-nat-test` when running tests.examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py (1)
83-87: Prefix unused variables with underscore.The static analysis correctly identifies that
ctxandlastare unpacked but never used. Prefix them with underscores to indicate they are intentionally ignored.Proposed fix
- result, ctx, last = initiate_group_chat( + result, _ctx, _last = initiate_group_chat( pattern=pattern, messages=config.task, max_rounds=config.max_rounds, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py` around lines 83 - 87, The tuple unpack from initiate_group_chat currently binds unused variables ctx and last; rename them to _ctx and _last (or _ , _) to signal they are intentionally ignored. Update the call site where initiate_group_chat is assigned (the line assigning result, ctx, last) to use result, _ctx, _last so static analysis no longer reports unused variables while preserving result usage.packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py (1)
46-50: Consider handling potential schema generation failures.If
fn.input_schema.model_json_schema()raises an exception (e.g., due to an unserializable type), the tool wrapper will fail. A similar pattern is tested inpackages/nvidia_nat_strands/tests/test_strands_tool_wrapper.py(lines 60-61).Proposed defensive handling
+ schema = None + if fn.input_schema: + try: + schema = fn.input_schema.model_json_schema() + except Exception: + schema = None + return Tool( name=name, description=fn.description or name, func_or_tool=call_function, - parameters_json_schema=( - fn.input_schema.model_json_schema() - if fn.input_schema - else None - ), + parameters_json_schema=schema, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py` around lines 46 - 50, The parameters_json_schema assignment can raise if fn.input_schema.model_json_schema() fails; update the code in tool_wrapper.py where parameters_json_schema is set (the expression using fn.input_schema.model_json_schema()) to defensively call model_json_schema() inside a try/except, catching exceptions (e.g., TypeError/ValueError/Exception), fallback to None on failure, and emit a warning or debug log including fn or function name to aid diagnosis; ensure the variable name parameters_json_schema and object fn/input_schema are referenced in the fix so reviewers can locate the change.examples/frameworks/nat_ag2_demo/pyproject.toml (1)
9-11: Package include pattern may exclude the base package.The pattern
nat_ag2_demo.*will only match submodules but not the basenat_ag2_demopackage itself. Consider including both the base package and its submodules.Proposed fix
[tool.setuptools.packages.find] where = ["src"] -include = ["nat_ag2_demo.*"] +include = ["nat_ag2_demo", "nat_ag2_demo.*"]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/pyproject.toml` around lines 9 - 11, The package include pattern currently only matches submodules (nat_ag2_demo.*) and will omit the base package; update the setuptools find configuration to include both the base package and its subpackages by adding the base package name alongside the glob pattern (e.g., include both "nat_ag2_demo" and "nat_ag2_demo.*") so that the top-level nat_ag2_demo package is packaged as well.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/frameworks/nat_ag2_demo/configs/config.yml`:
- Around line 1-12: Add the standard SPDX Apache-2.0 header at the top of this
YAML file so it begins with the required license statement; update the file
containing the "functions" block (with function name ag2_research and _type
ag2_research_team) and the "llm" block (with name default and _type nim) by
inserting the SPDX Apache-2.0 header template as the very first lines of the
file before any YAML content.
In `@examples/frameworks/nat_ag2_demo/pyproject.toml`:
- Around line 1-19: Add the SPDX license header to the top of this
pyproject.toml (the same SPDX-License-Identifier: Apache-2.0 line used in the
root pyproject.toml); open the file containing the [build-system] section and
the project entry with name = "nat-ag2-demo" and insert the SPDX header as the
very first line so the file begins with the license identifier before any TOML
tables or keys.
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py`:
- Around line 1-17: Add the required SPDX Apache-2.0 header at the top of the
file and remove the unused import of BaseModel; specifically, insert the
standard SPDX Apache-2.0 license header as the first lines of this module and
delete the "from pydantic import BaseModel" import since AG2ResearchConfig
already extends FunctionBaseConfig (see AG2ResearchConfig and FunctionBaseConfig
references) to eliminate the unused symbol.
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/register.py`:
- Line 1: Add the standard SPDX Apache-2.0 license header at the top of this
source file; currently the file only contains the import statement "from .
import ag2_research_team # noqa: F401", so prepend the required SPDX Apache-2.0
header template as the very first lines of register.py (above the import) to
satisfy the project licensing guideline.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/__init__.py`:
- Line 1: Add the required SPDX Apache-2.0 file header to the top of this module
(packages/nvidia_nat_ag2/src/nat/plugins/ag2/__init__.py) above the module
docstring: insert the standard SPDX-License-Identifier: Apache-2.0 line plus the
full Apache-2.0 license header block including the updated copyright year and
copyright owner, ensuring the header appears before the existing module
docstring in __init__.py so repository pre-commit checks pass.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 1-5: Add the standard SPDX Apache-2.0 license header to the top of
this module (before the module docstring) in
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py so the file
begins with the required license template; ensure the header appears above the
existing module docstring (the """Profiler callback handler for AG2...""" block)
and matches the project's SPDX Apache-2.0 format exactly.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py`:
- Around line 1-6: Add the standard SPDX Apache-2.0 license header at the very
top of the file packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py (above the
module docstring triple quotes) so the file begins with the required SPDX
boilerplate; ensure the header matches the project’s standard template exactly
(SPDX-License-Identifier: Apache-2.0) and keep the existing module docstring and
LLM wrapper code (LLMConfig-related content) unchanged below it.
- Around line 28-32: The api_key fields in the three AG2 provider wrappers are
passing SecretStr objects instead of raw strings; import get_secret_value from
nat.data_models.common and replace the direct llm_config.api_key usage with
get_secret_value(llm_config.api_key) in openai_ag2 (config_dict["api_key"]),
nim_ag2 (config_dict["api_key"]) and azure_ag2 (config_dict["api_key"]) so the
actual secret string is used for authentication.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py`:
- Around line 1-13: Add the missing SPDX Apache-2.0 license header at the very
top of this file (before the module docstring) in
packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py; ensure the standard
single-line SPDX identifier "SPDX-License-Identifier: Apache-2.0" is the first
line of the file so the header applies to the entire module that contains the
top-level docstring and imports (references: the module itself tool_wrapper.py
and symbols like Builder, Function, register_tool_wrapper, Tool).
In `@packages/nvidia_nat_ag2/tests/test_import.py`:
- Around line 1-13: This file is missing the required SPDX Apache-2.0 header;
add the standard SPDX Apache-2.0 license header as the first line(s) of
packages/nvidia_nat_ag2/tests/test_import.py above the existing module docstring
so the file begins with the SPDX header, preserving the rest of the file
including the test_register_module_imports() and test_framework_enum_has_ag2()
functions and their imports.
---
Nitpick comments:
In `@docs/source/components/integrations/frameworks.md`:
- Around line 93-94: The paragraph contains three successive sentences that all
start with "AG2", hurting readability; rewrite one or two of those sentences to
vary sentence openings (for example start one with "The framework" or "It
provides") and rephrase to keep the same meaning about community maintenance,
multi-agent orchestration, and features (GroupChat patterns, tool integration,
flexible LLM configuration and MCP integration) while ensuring the term "AG2"
still appears at least once for clarity.
In `@examples/frameworks/nat_ag2_demo/pyproject.toml`:
- Around line 9-11: The package include pattern currently only matches
submodules (nat_ag2_demo.*) and will omit the base package; update the
setuptools find configuration to include both the base package and its
subpackages by adding the base package name alongside the glob pattern (e.g.,
include both "nat_ag2_demo" and "nat_ag2_demo.*") so that the top-level
nat_ag2_demo package is packaged as well.
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py`:
- Around line 83-87: The tuple unpack from initiate_group_chat currently binds
unused variables ctx and last; rename them to _ctx and _last (or _ , _) to
signal they are intentionally ignored. Update the call site where
initiate_group_chat is assigned (the line assigning result, ctx, last) to use
result, _ctx, _last so static analysis no longer reports unused variables while
preserving result usage.
In `@packages/nvidia_nat_ag2/pyproject.toml`:
- Around line 14-18: Add `nvidia-nat-test` as an optional test dependency by
defining a `test` extra and including the package there (e.g., add a
`project.optional-dependencies` / `tool.setuptools_dynamic_dependencies` entry
for the `test` extra containing "nvidia-nat-test >= {version}" or the
appropriate version constraint); update the existing dependencies block
(referenced by the `dependencies` symbol) to keep runtime deps unchanged and add
the new `test` extra so CI and local test tooling can install `nvidia-nat-test`
when running tests.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py`:
- Around line 46-50: The parameters_json_schema assignment can raise if
fn.input_schema.model_json_schema() fails; update the code in tool_wrapper.py
where parameters_json_schema is set (the expression using
fn.input_schema.model_json_schema()) to defensively call model_json_schema()
inside a try/except, catching exceptions (e.g., TypeError/ValueError/Exception),
fallback to None on failure, and emit a warning or debug log including fn or
function name to aid diagnosis; ensure the variable name parameters_json_schema
and object fn/input_schema are referenced in the fix so reviewers can locate the
change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f8baf37a-70e4-494f-9875-cae292d201ce
📒 Files selected for processing (17)
README.mddocs/source/components/integrations/frameworks.mdexamples/frameworks/nat_ag2_demo/configs/config.ymlexamples/frameworks/nat_ag2_demo/pyproject.tomlexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/__init__.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/register.pypackages/nvidia_nat_ag2/pyproject.tomlpackages/nvidia_nat_ag2/src/nat/plugins/ag2/__init__.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/register.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.pypackages/nvidia_nat_ag2/tests/__init__.pypackages/nvidia_nat_ag2/tests/test_import.pypackages/nvidia_nat_core/src/nat/builder/framework_enum.pypyproject.toml
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (1)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py (1)
46-47:⚠️ Potential issue | 🔴 CriticalUnwrap
OptionalSecretStrbefore building AG2 provider config.All three wrappers pass
llm_config.api_keydirectly. These configs use secret-typed fields, so you should pass the raw string value instead.🔧 Proposed fix
from collections.abc import AsyncGenerator from autogen import LLMConfig from nat.builder import Builder from nat.builder.framework_enum import LLMFrameworkEnum from nat.cli.register_workflow import register_llm_client +from nat.data_models.common import get_secret_value from nat.llm.azure_openai_llm import AzureOpenAIModelConfig from nat.llm.nim_llm import NIMModelConfig from nat.llm.openai_llm import OpenAIModelConfig @@ config_dict = { "api_type": "openai", "model": llm_config.model, - "api_key": llm_config.api_key, + "api_key": get_secret_value(llm_config.api_key), } @@ config_dict = { "api_type": "openai", "model": llm_config.model, - "api_key": llm_config.api_key, + "api_key": get_secret_value(llm_config.api_key), "base_url": ( llm_config.api_base or "https://integrate.api.nvidia.com/v1" ), @@ config_dict = { "api_type": "azure", "model": llm_config.model, - "api_key": llm_config.api_key, + "api_key": get_secret_value(llm_config.api_key), "base_url": llm_config.api_base, "api_version": llm_config.api_version, }Also applies to: 68-69, 96-97
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py` around lines 46 - 47, The AG2 provider config is being built using llm_config.api_key directly; since api_key is an OptionalSecretStr you must unwrap it before use (call llm_config.api_key.get_secret_value() when not None) in each provider construction site (the three places where the dict/key "api_key" is set). Update the config builders to check for None and pass the raw string (or None) instead of the SecretStr wrapper so the AG2 client receives a plain string.
🧹 Nitpick comments (2)
examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py (1)
46-49: Add a minimum bound formax_rounds.
max_roundsshould reject non-positive values to prevent invalid chat loop settings.🔧 Proposed fix
max_rounds: int = Field( default=10, + ge=1, description="Max GroupChat rounds.", )As per coding guidelines: "Validate and sanitise all user input, especially in web or CLI interfaces."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py` around lines 46 - 49, The max_rounds Field currently allows non-positive values; update its validation in the model (the place where max_rounds: int = Field(...) is defined) to enforce a minimum of 1 (e.g., add ge=1 to the Field call) or add a Pydantic validator on max_rounds to raise a ValueError for values <= 0 so non-positive inputs are rejected and the chat loop cannot be configured with invalid rounds.packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py (1)
39-42: Add Google-style docstrings to each public wrapper function.The module has docs, but each public wrapper should also include a Google-style docstring (
Args,Yields/Returns).As per coding guidelines: "Provide Google-style docstrings for every public module, class, function and CLI command."
Also applies to: 61-64, 85-88
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py` around lines 39 - 42, Add Google-style docstrings to the public wrapper functions in this module (e.g., openai_ag2 and the other public wrappers referenced in the review) describing parameters and return/yield values: include an Args section that documents llm_config: OpenAIModelConfig and _builder: Builder, and a Yields (or Returns) section that documents the AsyncGenerator[LLMConfig, None] result; place the docstring immediately below each function signature using the Google style with brief one-line summary, Args, and Yields/Returns entries.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py`:
- Line 98: The call to initiate_group_chat unpacks three values into result,
ctx, last but ctx and last are unused and trigger RUF059; fix this by changing
the unpacking to only capture the used value(s) (e.g., assign unused values to _
or use a throwaway tuple like result, _ctx, _last = initiate_group_chat(...) or
result, *_ = initiate_group_chat(...)) so that initiate_group_chat remains
invoked but ruff no longer flags unused variables; update any references to
ctx/last accordingly if they are actually needed elsewhere.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 58-61: The patch currently binds a single handler instance so only
the first AG2ProfilerHandler receives events; change patch()/unpatch() to
maintain a collection (e.g., _handlers set/list) instead of a single
closure-bound handler: on patch() add the current handler instance to
AG2ProfilerHandler._handlers and only install the wrapped create once (using
_original_create/_patch_count), and in the installed wrapper iterate over
AG2ProfilerHandler._handlers to dispatch events to every registered handler;
similarly on unpatch() remove the handler from _handlers and only restore the
original create when _handlers becomes empty (update the places around
AG2ProfilerHandler._patch_lock/_patch_count/_original_create and the dispatching
logic referenced at lines ~67-70 and ~141 to use the handlers collection).
- Around line 82-83: The code is currently emitting raw kwargs/result/error via
StreamEventData(input=str(kwargs)) and similar UsageInfo fields; replace these
with a sanitized representation by introducing and calling a sanitizer (e.g.,
sanitize_event_payload(payload)) that returns a redacted dict/string before
emitting to StreamEventData or UsageInfo. Update every emission site that uses
str(kwargs), str(result), or str(error) (references: StreamEventData input and
UsageInfo payloads) to pass the sanitized output, and ensure the sanitizer
strips or masks common sensitive keys (password, token, secret, api_key, auth,
credentials) and handles nested structures safely.
- Around line 145-150: The unpatch path is decrementing
AG2ProfilerHandler._patch_count even when it's already zero, allowing underflow;
update the unpatch logic in AG2ProfilerHandler so you first check if
AG2ProfilerHandler._patch_count is > 0 before decrementing (or clamp it after
decrement), and only proceed to early-return logic when appropriate; keep the
existing guards that check AG2ProfilerHandler._original_create but ensure the
decrement is conditional to prevent negative refcounts and preserve correct
patch/unpatch lifecycle behavior.
---
Duplicate comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py`:
- Around line 46-47: The AG2 provider config is being built using
llm_config.api_key directly; since api_key is an OptionalSecretStr you must
unwrap it before use (call llm_config.api_key.get_secret_value() when not None)
in each provider construction site (the three places where the dict/key
"api_key" is set). Update the config builders to check for None and pass the raw
string (or None) instead of the SecretStr wrapper so the AG2 client receives a
plain string.
---
Nitpick comments:
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py`:
- Around line 46-49: The max_rounds Field currently allows non-positive values;
update its validation in the model (the place where max_rounds: int = Field(...)
is defined) to enforce a minimum of 1 (e.g., add ge=1 to the Field call) or add
a Pydantic validator on max_rounds to raise a ValueError for values <= 0 so
non-positive inputs are rejected and the chat loop cannot be configured with
invalid rounds.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py`:
- Around line 39-42: Add Google-style docstrings to the public wrapper functions
in this module (e.g., openai_ag2 and the other public wrappers referenced in the
review) describing parameters and return/yield values: include an Args section
that documents llm_config: OpenAIModelConfig and _builder: Builder, and a Yields
(or Returns) section that documents the AsyncGenerator[LLMConfig, None] result;
place the docstring immediately below each function signature using the Google
style with brief one-line summary, Args, and Yields/Returns entries.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 5f335330-49db-47ce-bbde-27ef482d660e
📒 Files selected for processing (13)
examples/frameworks/nat_ag2_demo/configs/config.ymlexamples/frameworks/nat_ag2_demo/pyproject.tomlexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/__init__.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/register.pypackages/nvidia_nat_ag2/pyproject.tomlpackages/nvidia_nat_ag2/src/nat/plugins/ag2/__init__.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/register.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.pypackages/nvidia_nat_ag2/tests/__init__.pypackages/nvidia_nat_ag2/tests/test_import.py
✅ Files skipped from review due to trivial changes (3)
- examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/init.py
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/init.py
- packages/nvidia_nat_ag2/tests/init.py
🚧 Files skipped from review as they are similar to previous changes (5)
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/register.py
- examples/frameworks/nat_ag2_demo/pyproject.toml
- packages/nvidia_nat_ag2/tests/test_import.py
- examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/register.py
- packages/nvidia_nat_ag2/pyproject.toml
examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (2)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py (2)
58-61:⚠️ Potential issue | 🟠 Major
patch_countimplies multi-handler support, but only the first handler gets events.
patch()increments_patch_count, but the closure binds a singlehandler=self; later handlers are never registered for dispatch.🔧 Suggested fix
class AG2ProfilerHandler(BaseProfilerCallback): _original_create = None _patch_lock = threading.Lock() _patch_count = 0 + _handlers: set["AG2ProfilerHandler"] = set() def patch(self) -> None: ... with AG2ProfilerHandler._patch_lock: AG2ProfilerHandler._patch_count += 1 + AG2ProfilerHandler._handlers.add(self) if AG2ProfilerHandler._original_create is not None: return # Already patched AG2ProfilerHandler._original_create = OpenAIWrapper.create + original_create = AG2ProfilerHandler._original_create - handler = self def patched_create(wrapper_self: Any, *args: Any, **kwargs: Any) -> Any: + handlers = tuple(AG2ProfilerHandler._handlers) ... - handler.step_manager.push_intermediate_step(start_payload) + for handler in handlers: + handler.step_manager.push_intermediate_step(start_payload) ... - result = AG2ProfilerHandler._original_create(wrapper_self, *args, **kwargs) + result = original_create(wrapper_self, *args, **kwargs) ...Also applies to: 67-70, 90-90
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py` around lines 58 - 61, The patch logic currently increments AG2ProfilerHandler._patch_count but binds a single handler via the closure so only the first instance receives events; change to maintain a collection (e.g., AG2ProfilerHandler._handlers = set() or list) and on patch() add self to that collection, and on unpatch()/close() remove self; replace the closure-bound single-callback (the function stored in AG2ProfilerHandler._original_create / _patched_create) so it iterates over AG2ProfilerHandler._handlers and dispatches events to each handler instance; ensure _patch_count and restoring of the original create function still happen (only restore when _patch_count reaches zero) and update references to _patched_create, _original_create, _patch_count, and the new _handlers set accordingly.
82-83:⚠️ Potential issue | 🟠 MajorAvoid emitting raw
kwargs/result/error into profiler events.
str(kwargs),str(result), andstr(e)can leak secrets and PII to the event stream.🔒 Suggested fix
+SENSITIVE_KEYS = {"password", "token", "secret", "api_key", "auth", "authorization", "credentials"} + +def _sanitize_payload(value: Any) -> Any: + if isinstance(value, dict): + return { + k: ("***REDACTED***" if str(k).lower() in SENSITIVE_KEYS else _sanitize_payload(v)) + for k, v in value.items() + } + if isinstance(value, (list, tuple)): + return [_sanitize_payload(v) for v in value] + return value ... - data=StreamEventData(input=str(kwargs)), + data=StreamEventData(input=_sanitize_payload(kwargs)), ... - data=StreamEventData(input=str(kwargs), output=str(result)), + data=StreamEventData( + input=_sanitize_payload(kwargs), + output=_sanitize_payload(result), + ), ... - data=StreamEventData(input=str(kwargs), output=str(e)), + data=StreamEventData( + input=_sanitize_payload(kwargs), + output={"error": type(e).__name__, "message": str(e)}, + ),As per coding guidelines: "Validate and sanitise all user input, especially in web or CLI interfaces" and "Check for security vulnerabilities and potential issues."
Also applies to: 106-108, 127-130
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py` around lines 82 - 83, The profiler events are currently emitting raw sensitive data via StreamEventData(input=str(kwargs)), StreamEventData(input=str(result)), and StreamEventData(input=str(e)); replace these direct serializations with a sanitization/summary step (e.g., implement and call a helper like sanitize_for_event(value) or summarize_for_event(value) that redacts secrets/PII, returns type, length, and a truncated non-sensitive preview) and use that helper in the three places where StreamEventData is constructed; update CallbackHandler (or the class/function containing those StreamEventData calls) to call the sanitizer and only include non-sensitive metadata in usage/profiler events.
🧹 Nitpick comments (1)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py (1)
18-20: Wrap code entities in backticks inside the module docstring.In the module docstring, wrap
LLMConfigin backticks to match docstring lint expectations.As per coding guidelines: "Surround code entities with backticks to avoid Vale false-positives in docstrings."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py` around lines 18 - 20, Update the module docstring in packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py to wrap the code entity LLMConfig in backticks (`LLMConfig`) wherever it appears (e.g., the sentence "AG2 uses LLMConfig (a configuration object) ..."). Ensure any other code entities in that docstring are similarly backticked to satisfy the docstring lint rule and avoid Vale false-positives.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 94-96: patched_create dereferences the class variable
AG2ProfilerHandler._original_create while unpatch() can set it to None
concurrently; capture AG2ProfilerHandler._original_create into a local variable
at the start of patched_create (e.g., local_original =
AG2ProfilerHandler._original_create), check that local_original is not None
(return/raise or no-op if it is), and then call local_original(wrapper_self,
*args, **kwargs) instead of dereferencing the class attribute directly. Apply
the same change to the other patched function(s) around the 145-159 region that
similarly use AG2ProfilerHandler._original_create so they also capture and check
the local copy before invoking it.
- Around line 18-20: Update the module/docstring text in callback_handler.py to
use the project API namespace casing `nat` instead of `NAT` — e.g., change the
sentence that reads "Patches AG2's OpenAIWrapper to capture LLM call timings and
token usage for NAT's profiling pipeline." to use "nat" (or `nat`) and scan the
module for other prose references to "NAT" (particularly around the AG2
OpenAIWrapper patch or any top-level docstring) and replace them with the
correct `nat` casing.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py`:
- Around line 40-43: Add concise Google-style docstrings to the three public
wrapper functions openai_ag2, nim_ag2, and azure_ag2: for each function document
its purpose (brief one-line summary), Args (llm_config: OpenAIModelConfig or
equivalent, _builder: Builder), Yields (LLMConfig AsyncGenerator), and any
important behavior/side-effects; place the docstring immediately under each
function definition (openai_ag2, nim_ag2, azure_ag2) using Google-style sections
(Args, Yields, Returns/Raises if applicable) and keep them short and focused.
---
Duplicate comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 58-61: The patch logic currently increments
AG2ProfilerHandler._patch_count but binds a single handler via the closure so
only the first instance receives events; change to maintain a collection (e.g.,
AG2ProfilerHandler._handlers = set() or list) and on patch() add self to that
collection, and on unpatch()/close() remove self; replace the closure-bound
single-callback (the function stored in AG2ProfilerHandler._original_create /
_patched_create) so it iterates over AG2ProfilerHandler._handlers and dispatches
events to each handler instance; ensure _patch_count and restoring of the
original create function still happen (only restore when _patch_count reaches
zero) and update references to _patched_create, _original_create, _patch_count,
and the new _handlers set accordingly.
- Around line 82-83: The profiler events are currently emitting raw sensitive
data via StreamEventData(input=str(kwargs)), StreamEventData(input=str(result)),
and StreamEventData(input=str(e)); replace these direct serializations with a
sanitization/summary step (e.g., implement and call a helper like
sanitize_for_event(value) or summarize_for_event(value) that redacts
secrets/PII, returns type, length, and a truncated non-sensitive preview) and
use that helper in the three places where StreamEventData is constructed; update
CallbackHandler (or the class/function containing those StreamEventData calls)
to call the sanitizer and only include non-sensitive metadata in usage/profiler
events.
---
Nitpick comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py`:
- Around line 18-20: Update the module docstring in
packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py to wrap the code entity
LLMConfig in backticks (`LLMConfig`) wherever it appears (e.g., the sentence
"AG2 uses LLMConfig (a configuration object) ..."). Ensure any other code
entities in that docstring are similarly backticked to satisfy the docstring
lint rule and avoid Vale false-positives.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 145e15dc-8b31-4868-bcbe-a8a314a7ce4c
📒 Files selected for processing (2)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py (2)
56-60: Add Args and Returns sections to the docstring.Per coding guidelines, public functions require Google-style docstrings. The current docstring is missing the
ArgsandReturnssections.📝 Proposed docstring improvement
async def ag2_research_team( config: AG2ResearchConfig, builder: Builder, ) -> str: - """Run a 2-agent research team with AG2.""" + """Run a 2-agent research team with AG2. + + Args: + config: Configuration specifying the research task, LLM config name, + and maximum chat rounds. + builder: NAT builder instance for resolving LLM configurations. + + Returns: + The last substantive message from the chat history, or + "Research complete." if no content is found. + """As per coding guidelines: "Provide Google-style docstrings for every public module, class, function and CLI command."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py` around lines 56 - 60, The docstring for the public function ag2_research_team is missing Google-style Args and Returns sections; update the triple-quoted docstring of ag2_research_team to include an Args section documenting the config: AG2ResearchConfig and builder: Builder parameters (types and brief descriptions) and a Returns section describing the returned str (what it represents), following the project’s Google-style docstring format.
98-102: Synchronous AG2 call in async context.
initiate_group_chatappears to be a synchronous/blocking call within an async function. For demo purposes this is acceptable, but in production scenarios this could block the event loop. Consider documenting this limitation or wrapping inasyncio.to_thread()if non-blocking behavior is needed.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py` around lines 98 - 102, The call to initiate_group_chat is a blocking/synchronous call inside an async function (using pattern, messages=config.task, max_rounds=config.max_rounds) which can block the event loop; either document this limitation or make the call non-blocking by running initiate_group_chat in a thread (e.g., wrap it with asyncio.to_thread) and await that result, ensuring the surrounding async function awaits the threaded call and preserves result, _ctx, _last semantics.packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py (1)
39-41: AddClassVarannotation to clarify class-level mutable attribute.The static analysis warning RUF012 flags the mutable default
[]on line 41. While this is intentionally a shared class variable, addingClassVarmakes the intent explicit and silences the linter.♻️ Proposed fix
+from typing import Any, ClassVar, Optional import threading import time -from typing import Any ... class AG2ProfilerHandler(BaseProfilerCallback): """Instruments AG2 agents for profiling.""" - _original_create = None - _patch_lock = threading.Lock() - _handlers: list["AG2ProfilerHandler"] = [] + _original_create: ClassVar[Optional[Any]] = None + _patch_lock: ClassVar[threading.Lock] = threading.Lock() + _handlers: ClassVar[list["AG2ProfilerHandler"]] = []🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py` around lines 39 - 41, The mutable class-level variable _handlers should be annotated with typing.ClassVar to make its intent explicit and silence RUF012; update the declaration of _handlers to use ClassVar (e.g., _handlers: ClassVar[list["AG2ProfilerHandler"]]) and add an import for ClassVar from typing if not already present, keeping the list initializer unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 67-72: The fallback in patched_create causes infinite recursion
because when AG2ProfilerHandler._original_create is None the code calls
wrapper_self.create which is still bound to patched_create; change the fallback
to not call wrapper_self.create: instead detect original is None and raise a
clear RuntimeError (or log an error and return/raise) with context, e.g. inside
patched_create check AG2ProfilerHandler._original_create and if None call
logger.error(...) and raise RuntimeError("AG2ProfilerHandler unpatched:
original_create missing") rather than calling wrapper_self.create; update
references in patched_create and ensure any tests that expect graceful handling
adjust accordingly.
---
Nitpick comments:
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py`:
- Around line 56-60: The docstring for the public function ag2_research_team is
missing Google-style Args and Returns sections; update the triple-quoted
docstring of ag2_research_team to include an Args section documenting the
config: AG2ResearchConfig and builder: Builder parameters (types and brief
descriptions) and a Returns section describing the returned str (what it
represents), following the project’s Google-style docstring format.
- Around line 98-102: The call to initiate_group_chat is a blocking/synchronous
call inside an async function (using pattern, messages=config.task,
max_rounds=config.max_rounds) which can block the event loop; either document
this limitation or make the call non-blocking by running initiate_group_chat in
a thread (e.g., wrap it with asyncio.to_thread) and await that result, ensuring
the surrounding async function awaits the threaded call and preserves result,
_ctx, _last semantics.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 39-41: The mutable class-level variable _handlers should be
annotated with typing.ClassVar to make its intent explicit and silence RUF012;
update the declaration of _handlers to use ClassVar (e.g., _handlers:
ClassVar[list["AG2ProfilerHandler"]]) and add an import for ClassVar from typing
if not already present, keeping the list initializer unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: de8daeac-d63c-49f4-8d1d-2031c1ae83f2
📒 Files selected for processing (3)
examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py
Outdated
Show resolved
Hide resolved
willkill07
added
feature request
willkill07
requested changes
Mar 18, 2026
Member
willkill07
left a comment
willkill07
left a comment
There was a problem hiding this comment.
Here's a first round of feedback on the PR.
Member
There was a problem hiding this comment.
This is not a complete configuration
examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py
Outdated
Show resolved
Hide resolved
jsun-m
force-pushed
the
vasiliyr/03152026
branch
from
74b7d15 to
0893dfd
Compare
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
jsun-m
force-pushed
the
vasiliyr/03152026
branch
from
0893dfd to
e39a534
Compare
# Conflicts: # examples/A2A/currency_agent_a2a/uv.lock # examples/A2A/math_assistant_a2a/uv.lock # examples/A2A/math_assistant_a2a_protected/uv.lock # examples/HITL/por_to_jiratickets/uv.lock # examples/HITL/simple_calculator_hitl/uv.lock # examples/MCP/kaggle_mcp/uv.lock # examples/MCP/service_account_auth_mcp/uv.lock # examples/MCP/simple_auth_mcp/uv.lock # examples/MCP/simple_calculator_fastmcp/uv.lock # examples/MCP/simple_calculator_fastmcp_protected/uv.lock # examples/MCP/simple_calculator_mcp/uv.lock # examples/MCP/simple_calculator_mcp_protected/uv.lock # examples/RAG/simple_rag/uv.lock # examples/advanced_agents/alert_triage_agent/uv.lock # examples/agents/uv.lock # examples/control_flow/hybrid_control_flow/uv.lock # examples/control_flow/parallel_executor/uv.lock # examples/control_flow/router_agent/uv.lock # examples/control_flow/sequential_executor/uv.lock # examples/custom_functions/automated_description_generation/uv.lock # examples/custom_functions/plot_charts/uv.lock # examples/documentation_guides/uv.lock # examples/documentation_guides/workflows/text_file_ingest/uv.lock # examples/dynamo_integration/latency_sensitivity_demo/uv.lock # examples/dynamo_integration/react_benchmark_agent/uv.lock # examples/evaluation_and_profiling/email_phishing_analyzer/uv.lock # examples/evaluation_and_profiling/simple_calculator_eval/uv.lock # examples/evaluation_and_profiling/simple_web_query_eval/uv.lock # examples/finetuning/dpo_tic_tac_toe/uv.lock # examples/finetuning/rl_with_openpipe_art/uv.lock # examples/frameworks/adk_demo/uv.lock # examples/frameworks/agno_personal_finance/uv.lock # examples/frameworks/haystack_deep_research_agent/uv.lock # examples/frameworks/multi_frameworks/uv.lock # examples/frameworks/nat_autogen_demo/uv.lock # examples/frameworks/semantic_kernel_demo/uv.lock # examples/frameworks/strands_demo/uv.lock # examples/front_ends/per_user_workflow/uv.lock # examples/front_ends/simple_auth/uv.lock # examples/front_ends/simple_calculator_custom_routes/uv.lock # examples/getting_started/simple_calculator/uv.lock # examples/getting_started/simple_web_query/uv.lock # examples/memory/redis/uv.lock # examples/notebooks/uv.lock # examples/object_store/user_report/uv.lock # examples/observability/simple_calculator_observability/uv.lock # examples/prompt_from_file/uv.lock # examples/safety_and_security/retail_agent/uv.lock # uv.lock
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
Author
|
@coderabbitai resume |
✅ Actions performedReviews resumed. |
coderabbitai
bot
added
the
DO NOT MERGE
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (1)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py (1)
65-87:⚠️ Potential issue | 🟠 MajorInstrumentation lifecycle is instance-local while patch target is global.
Multiple
AG2ProfilerHandlerinstances can overwrite each other’s patch state and restore order, causing lost instrumentation or incorrect unpatch behavior. This needs class-level coordination (shared lock + shared original + refcount/handler registry).Also applies to: 88-105
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py` around lines 65 - 87, The instrumentation state is currently instance-local causing race/overwrite between AG2ProfilerHandler instances; change the lifecycle to class-level coordination by adding class-level vars (e.g., AG2ProfilerHandler._lock, _refcount, _original_create, _instrumented) and protect instrument() and uninstrument() with that lock: on instrument() increment _refcount and only patch OpenAIWrapper.create (store original in _original_create) when _refcount transitions 0→1, and on uninstrument() decrement _refcount and only restore OpenAIWrapper.create when _refcount transitions 1→0; ensure thread-safety around these operations and keep existing debug logs (referencing methods instrument(), uninstrument(), and symbol OpenAIWrapper.create).
🧹 Nitpick comments (6)
packages/nvidia_nat_ag2/pyproject.toml (1)
43-45: Addnvidia-nat-testas an optional test dependency.
nvidia-nat-testis declared in[tool.uv.sources]but not listed as a dependency. Per coding guidelines, it should be declared as an optional dependency in thetestextra.♻️ Proposed fix to add optional test dependency
[tool.setuptools_dynamic_dependencies] dependencies = [ "nvidia-nat-core == {version}", "ag2[openai] ~= 0.11", ] +[tool.setuptools_dynamic_dependencies.optional-dependencies] +test = [ + "nvidia-nat-test == {version}", +] + [project.entry-points.'nat.components'] nat_ag2 = "nat.plugins.ag2.register"As per coding guidelines: "
nvidia-nat-testshould likely be listed as an optional dependency in thetestextra".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/pyproject.toml` around lines 43 - 45, The pyproject.toml declares nvidia-nat-test in [tool.uv.sources] but doesn't add it to the test extra; update pyproject.toml to include nvidia-nat-test as an optional dependency under the test extra so consumers can opt into it. Locate the [tool.uv.sources] entry for "nvidia-nat-test" and add the package name "nvidia-nat-test" to the test extra (project optional-dependencies / tool.uv.extras "test") so the source mapping is usable as an optional test dependency.packages/nvidia_nat_ag2/tests/test_callback_handler.py (1)
386-387: Consider removing the manualpytest.main()entrypoint.For repository-managed pytest runs, this block is redundant and can be dropped to keep tests lean.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/tests/test_callback_handler.py` around lines 386 - 387, Remove the manual test entrypoint block that calls pytest.main — delete the "if __name__ == \"__main__\":" conditional and its pytest.main([__file__, "-v"]) call so the test module relies on repository-managed pytest runs; ensure no other code depends on that runtime invocation and leave the test functions/classes (and imports) intact.examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_async_team.py (1)
131-131: Consider addinginput_schemafor better LLM tool parameter discovery.Same pattern as other workflow files—
FunctionInfo.from_fn(_ag2_async_team_workflow)lacks an explicitinput_schema, which means the AG2 tool wrapper won't expose structured parameter information foruser_input.🔧 Proposed fix
+from pydantic import BaseModel + +class AsyncTeamInput(BaseModel): + """Input schema for the async team workflow.""" + user_input: str = Field(description="The user's question about LA traffic.") + ... - yield FunctionInfo.from_fn(_ag2_async_team_workflow) + yield FunctionInfo.from_fn(_ag2_async_team_workflow, input_schema=AsyncTeamInput)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_async_team.py` at line 131, The FunctionInfo.from_fn call for _ag2_async_team_workflow doesn't provide an input_schema so the AG2 tool wrapper can't expose structured parameters; update the yield to include an explicit input_schema describing the workflow's input (e.g., a schema for "user_input" with type string and an optional description/required flags) by passing input_schema=... to FunctionInfo.from_fn so downstream tools can discover the parameter shape for _ag2_async_team_workflow.examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_team.py (1)
36-137: Significant code duplication withag2_async_team.py.This file is nearly identical to
ag2_async_team.py—both usea_initiate_group_chat, create the same agent structure, and have identical chat history scanning logic. The only differences are the config/function names. Consider consolidating into a single implementation with a config flag, or clearly document why separate implementations are needed.💡 Consolidation suggestion
If both variants truly need to exist, consider extracting common logic into a shared helper:
# In a shared module (e.g., ag2_common.py) async def run_ag2_group_chat( llm_config, tools, query_agent_name: str, query_agent_instructions: str, final_agent_name: str, final_agent_instructions: str, user_input: str, max_rounds: int, approval_keyword: str = "APPROVE", ) -> str: """Shared AG2 group chat execution logic.""" # ... common implementation🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_team.py` around lines 36 - 137, The ag2_team implementation duplicates logic from ag2_async_team.py; extract the shared group-chat logic into a helper and have both register functions call it. Create a shared async helper (e.g., run_ag2_group_chat) that accepts llm_config, tools, query_agent_name, query_agent_instructions, final_agent_name, final_agent_instructions, user_input, max_rounds and approval_keyword and contains the ConversableAgent creation, tool registration, AutoPattern construction, a_initiate_group_chat call, and post-chat history scanning; then update ag2_team's _ag2_team_workflow to call that helper and return its result (similarly update ag2_async_team to reuse it), keeping AG2TeamConfig and ag2_team function signatures unchanged.examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py (2)
46-111: Consider adding error handling for robustness.The function lacks try/except handling around the AG2 workflow execution. Unlike
ag2_team.pyandag2_async_team.py, this implementation doesn't catch exceptions duringa_initiate_group_chator agent creation, which could cause unhandled failures.🛡️ Proposed error handling
async def _ag2_research_team(task: str) -> str: """Run the AG2 research team on the given task.""" + try: researcher = ConversableAgent( ... ) ... result, _ctx, _last = await a_initiate_group_chat(...) for msg in reversed(result.chat_history): ... return "Research complete." + except Exception: + logger.exception("Error in AG2 research team workflow") + return "An internal error occurred during the research workflow." yield FunctionInfo.from_fn(_ag2_research_team)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py` around lines 46 - 111, Wrap the AG2 workflow in ag2_research_team's inner _ag2_research_team with a try/except to catch exceptions from ConversableAgent creation, AutoPattern construction, and the a_initiate_group_chat call; on exception, log the error (including exception details) and either return a clear error string or re-raise a wrapped exception so the caller can handle it. Specifically, add a try around the blocks that instantiate ConversableAgent (researcher, writer, user), build AutoPattern, and await a_initiate_group_chat, catch Exception as e, use an available logger or builder/process logger to record f"ag2 research team error: {e}" and stack info, then return a failure message (e.g., "Research failed: <brief>") or rethrow after logging depending on surrounding conventions.
113-113: Consider addinginput_schemafor better LLM tool parameter discovery.Similar to the traffic status tool,
FunctionInfo.from_fn(_ag2_research_team)lacks an explicitinput_schema. If this function is exposed as a tool, the LLM won't receive structured parameter information fortask.🔧 Proposed fix
+from pydantic import BaseModel + +class ResearchTaskInput(BaseModel): + """Input schema for the research team.""" + task: str = Field(description="The research topic or question to investigate.") + ... - yield FunctionInfo.from_fn(_ag2_research_team) + yield FunctionInfo.from_fn(_ag2_research_team, input_schema=ResearchTaskInput)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py` at line 113, FunctionInfo.from_fn(_ag2_research_team) is missing an explicit input_schema so the LLM won't get structured parameter info for the task argument; update the registration to include an input_schema describing the task parameter (e.g., required string field "task" or a richer JSON/Pydantic schema) when calling FunctionInfo.from_fn or FunctionInfo(..., input_schema=...), referencing the _ag2_research_team function and the FunctionInfo.from_fn/FunctionInfo constructor to provide the schema for tool parameter discovery.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/source/components/integrations/frameworks.md`:
- Line 103: Replace the possessive phrasing "Fully supported through AG2's
`Tool` interface" in the frameworks.md AG2 capability row with non-possessive
wording such as "Fully supported through the `Tool` interface of AG2"; locate
the table cell under the "Tool Calling" row and update that exact string (the
phrase containing AG2's `Tool` interface) to the new wording.
In `@examples/frameworks/nat_ag2_demo/pyproject.toml`:
- Around line 24-26: The package discovery configuration under
[tool.setuptools.packages.find] incorrectly uses include = ["nat_ag2_demo.*"],
which excludes the base package and causes import/entrypoint failures; remove
the include key so setuptools can auto-discover packages (leave where = ["src"]
and delete include) so the base package nat_ag2_demo and its entry point
nat_ag2_demo.register are found.
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/traffic_status_tool.py`:
- Line 143: The tool registration omits an explicit input_schema so
FunctionInfo.from_fn(_traffic_status, ...) leaves parameters_json_schema=None
and the LLM won't get structured params for hwy and hour; fix by defining a
Pydantic model (e.g., class TrafficStatusInput(BaseModel): hwy: str; hour: int)
and pass it into FunctionInfo.from_fn via input_schema=TrafficStatusInput to
expose typed parameters for _traffic_status to AG2. Ensure the model field names
match the _traffic_status signature and update the FunctionInfo.from_fn call to
include input_schema=TrafficStatusInput.
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 130-133: The message_count extraction can raise TypeError for None
or non-sized values; replace len(kwargs.get("messages", [])) with a defensive
computation: pull messages = kwargs.get("messages"), then compute message_count
using a try/except TypeError (or check hasattr(messages, "__len__")) and default
to 0 when messages is None or not sized, then use that safe message_count in the
returned dict (refer to the return block in callback_handler.py that builds the
dict with "model" and "message_count").
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Line 110: Remove the unnecessary `@pytest.mark.asyncio` decorators from the
async tests in this module: search for occurrences of the symbol
"@pytest.mark.asyncio" in the test_tool_wrapper tests and delete those decorator
lines (keeping the async def test_* functions intact) so pytest-asyncio's
auto-detection is used instead.
---
Duplicate comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 65-87: The instrumentation state is currently instance-local
causing race/overwrite between AG2ProfilerHandler instances; change the
lifecycle to class-level coordination by adding class-level vars (e.g.,
AG2ProfilerHandler._lock, _refcount, _original_create, _instrumented) and
protect instrument() and uninstrument() with that lock: on instrument()
increment _refcount and only patch OpenAIWrapper.create (store original in
_original_create) when _refcount transitions 0→1, and on uninstrument()
decrement _refcount and only restore OpenAIWrapper.create when _refcount
transitions 1→0; ensure thread-safety around these operations and keep existing
debug logs (referencing methods instrument(), uninstrument(), and symbol
OpenAIWrapper.create).
---
Nitpick comments:
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_async_team.py`:
- Line 131: The FunctionInfo.from_fn call for _ag2_async_team_workflow doesn't
provide an input_schema so the AG2 tool wrapper can't expose structured
parameters; update the yield to include an explicit input_schema describing the
workflow's input (e.g., a schema for "user_input" with type string and an
optional description/required flags) by passing input_schema=... to
FunctionInfo.from_fn so downstream tools can discover the parameter shape for
_ag2_async_team_workflow.
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.py`:
- Around line 46-111: Wrap the AG2 workflow in ag2_research_team's inner
_ag2_research_team with a try/except to catch exceptions from ConversableAgent
creation, AutoPattern construction, and the a_initiate_group_chat call; on
exception, log the error (including exception details) and either return a clear
error string or re-raise a wrapped exception so the caller can handle it.
Specifically, add a try around the blocks that instantiate ConversableAgent
(researcher, writer, user), build AutoPattern, and await a_initiate_group_chat,
catch Exception as e, use an available logger or builder/process logger to
record f"ag2 research team error: {e}" and stack info, then return a failure
message (e.g., "Research failed: <brief>") or rethrow after logging depending on
surrounding conventions.
- Line 113: FunctionInfo.from_fn(_ag2_research_team) is missing an explicit
input_schema so the LLM won't get structured parameter info for the task
argument; update the registration to include an input_schema describing the task
parameter (e.g., required string field "task" or a richer JSON/Pydantic schema)
when calling FunctionInfo.from_fn or FunctionInfo(..., input_schema=...),
referencing the _ag2_research_team function and the
FunctionInfo.from_fn/FunctionInfo constructor to provide the schema for tool
parameter discovery.
In `@examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_team.py`:
- Around line 36-137: The ag2_team implementation duplicates logic from
ag2_async_team.py; extract the shared group-chat logic into a helper and have
both register functions call it. Create a shared async helper (e.g.,
run_ag2_group_chat) that accepts llm_config, tools, query_agent_name,
query_agent_instructions, final_agent_name, final_agent_instructions,
user_input, max_rounds and approval_keyword and contains the ConversableAgent
creation, tool registration, AutoPattern construction, a_initiate_group_chat
call, and post-chat history scanning; then update ag2_team's _ag2_team_workflow
to call that helper and return its result (similarly update ag2_async_team to
reuse it), keeping AG2TeamConfig and ag2_team function signatures unchanged.
In `@packages/nvidia_nat_ag2/pyproject.toml`:
- Around line 43-45: The pyproject.toml declares nvidia-nat-test in
[tool.uv.sources] but doesn't add it to the test extra; update pyproject.toml to
include nvidia-nat-test as an optional dependency under the test extra so
consumers can opt into it. Locate the [tool.uv.sources] entry for
"nvidia-nat-test" and add the package name "nvidia-nat-test" to the test extra
(project optional-dependencies / tool.uv.extras "test") so the source mapping is
usable as an optional test dependency.
In `@packages/nvidia_nat_ag2/tests/test_callback_handler.py`:
- Around line 386-387: Remove the manual test entrypoint block that calls
pytest.main — delete the "if __name__ == \"__main__\":" conditional and its
pytest.main([__file__, "-v"]) call so the test module relies on
repository-managed pytest runs; ensure no other code depends on that runtime
invocation and leave the test functions/classes (and imports) intact.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 66df5d85-4a9c-41f2-9f39-9f947ce1c80f
⛔ Files ignored due to path filters (50)
examples/A2A/currency_agent_a2a/uv.lockis excluded by!**/*.lockexamples/A2A/math_assistant_a2a/uv.lockis excluded by!**/*.lockexamples/A2A/math_assistant_a2a_protected/uv.lockis excluded by!**/*.lockexamples/HITL/por_to_jiratickets/uv.lockis excluded by!**/*.lockexamples/HITL/simple_calculator_hitl/uv.lockis excluded by!**/*.lockexamples/MCP/kaggle_mcp/uv.lockis excluded by!**/*.lockexamples/MCP/service_account_auth_mcp/uv.lockis excluded by!**/*.lockexamples/MCP/simple_auth_mcp/uv.lockis excluded by!**/*.lockexamples/MCP/simple_calculator_fastmcp/uv.lockis excluded by!**/*.lockexamples/MCP/simple_calculator_fastmcp_protected/uv.lockis excluded by!**/*.lockexamples/MCP/simple_calculator_mcp/uv.lockis excluded by!**/*.lockexamples/MCP/simple_calculator_mcp_protected/uv.lockis excluded by!**/*.lockexamples/RAG/simple_rag/uv.lockis excluded by!**/*.lockexamples/advanced_agents/alert_triage_agent/uv.lockis excluded by!**/*.lockexamples/agents/uv.lockis excluded by!**/*.lockexamples/control_flow/hybrid_control_flow/uv.lockis excluded by!**/*.lockexamples/control_flow/parallel_executor/uv.lockis excluded by!**/*.lockexamples/control_flow/router_agent/uv.lockis excluded by!**/*.lockexamples/control_flow/sequential_executor/uv.lockis excluded by!**/*.lockexamples/custom_functions/automated_description_generation/uv.lockis excluded by!**/*.lockexamples/custom_functions/plot_charts/uv.lockis excluded by!**/*.lockexamples/documentation_guides/uv.lockis excluded by!**/*.lockexamples/documentation_guides/workflows/text_file_ingest/uv.lockis excluded by!**/*.lockexamples/dynamo_integration/latency_sensitivity_demo/uv.lockis excluded by!**/*.lockexamples/dynamo_integration/react_benchmark_agent/uv.lockis excluded by!**/*.lockexamples/evaluation_and_profiling/email_phishing_analyzer/uv.lockis excluded by!**/*.lockexamples/evaluation_and_profiling/simple_calculator_eval/uv.lockis excluded by!**/*.lockexamples/evaluation_and_profiling/simple_web_query_eval/uv.lockis excluded by!**/*.lockexamples/finetuning/dpo_tic_tac_toe/uv.lockis excluded by!**/*.lockexamples/finetuning/rl_with_openpipe_art/uv.lockis excluded by!**/*.lockexamples/frameworks/adk_demo/uv.lockis excluded by!**/*.lockexamples/frameworks/agno_personal_finance/uv.lockis excluded by!**/*.lockexamples/frameworks/haystack_deep_research_agent/uv.lockis excluded by!**/*.lockexamples/frameworks/multi_frameworks/uv.lockis excluded by!**/*.lockexamples/frameworks/nat_ag2_demo/uv.lockis excluded by!**/*.lockexamples/frameworks/nat_autogen_demo/uv.lockis excluded by!**/*.lockexamples/frameworks/semantic_kernel_demo/uv.lockis excluded by!**/*.lockexamples/frameworks/strands_demo/uv.lockis excluded by!**/*.lockexamples/front_ends/per_user_workflow/uv.lockis excluded by!**/*.lockexamples/front_ends/simple_auth/uv.lockis excluded by!**/*.lockexamples/front_ends/simple_calculator_custom_routes/uv.lockis excluded by!**/*.lockexamples/getting_started/simple_calculator/uv.lockis excluded by!**/*.lockexamples/getting_started/simple_web_query/uv.lockis excluded by!**/*.lockexamples/memory/redis/uv.lockis excluded by!**/*.lockexamples/notebooks/uv.lockis excluded by!**/*.lockexamples/object_store/user_report/uv.lockis excluded by!**/*.lockexamples/observability/simple_calculator_observability/uv.lockis excluded by!**/*.lockexamples/prompt_from_file/uv.lockis excluded by!**/*.lockexamples/safety_and_security/retail_agent/uv.lockis excluded by!**/*.lockpackages/nvidia_nat_ag2/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (26)
README.mddocs/source/components/integrations/frameworks.mdexamples/frameworks/nat_ag2_demo/README.mdexamples/frameworks/nat_ag2_demo/configs/config-async.ymlexamples/frameworks/nat_ag2_demo/configs/config-eval.ymlexamples/frameworks/nat_ag2_demo/configs/config-research.ymlexamples/frameworks/nat_ag2_demo/configs/config.ymlexamples/frameworks/nat_ag2_demo/data/toy_data.jsonexamples/frameworks/nat_ag2_demo/pyproject.tomlexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/__init__.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_async_team.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_research_team.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/ag2_team.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/register.pyexamples/frameworks/nat_ag2_demo/src/nat_ag2_demo/traffic_status_tool.pypackages/nvidia_nat_ag2/pyproject.tomlpackages/nvidia_nat_ag2/src/nat/plugins/ag2/__init__.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/register.pypackages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.pypackages/nvidia_nat_ag2/tests/__init__.pypackages/nvidia_nat_ag2/tests/test_callback_handler.pypackages/nvidia_nat_ag2/tests/test_import.pypackages/nvidia_nat_ag2/tests/test_tool_wrapper.pypackages/nvidia_nat_core/src/nat/builder/framework_enum.py
✅ Files skipped from review due to trivial changes (8)
- packages/nvidia_nat_ag2/tests/init.py
- README.md
- examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/init.py
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/init.py
- examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/register.py
- packages/nvidia_nat_core/src/nat/builder/framework_enum.py
- packages/nvidia_nat_ag2/tests/test_import.py
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/register.py
🚧 Files skipped from review as they are similar to previous changes (6)
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py
- examples/frameworks/nat_ag2_demo/configs/config-research.yml
- examples/frameworks/nat_ag2_demo/README.md
- examples/frameworks/nat_ag2_demo/configs/config-eval.yml
- examples/frameworks/nat_ag2_demo/configs/config.yml
- packages/nvidia_nat_ag2/src/nat/plugins/ag2/llm.py
examples/frameworks/nat_ag2_demo/src/nat_ag2_demo/traffic_status_tool.py
Show resolved
Hide resolved
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
coderabbitai
bot
removed
the
DO NOT MERGE
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py (1)
81-83:⚠️ Potential issue | 🔴 CriticalShared monkey-patch lifecycle is still unsafe across multiple handler instances.
At Line 81 and Line 97, patch state is instance-scoped (
self._original_create), so patch/unpatch ordering across twoAG2ProfilerHandlerobjects can disable active instrumentation or restore a stale wrapper.Use class-level patch state (lock + refcount/handler registry) and restore the original method only when the last active handler unpatches.
Verification (read-only): expect to see only instance-level lifecycle fields and no class-level patch registry.
#!/bin/bash set -euo pipefail rg -n --type=py -C2 "class AG2ProfilerHandler|self\._original_create|self\._instrumented|OpenAIWrapper\.create\s*=" rg -n --type=py -C2 "_patch_count|_patch_lock|_handlers"Also applies to: 94-99
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py` around lines 81 - 83, The current monkey-patch in AG2ProfilerHandler stores patch state per-instance (self._original_create, self._instrumented) and directly assigns OpenAIWrapper.create, which leads to unsafe interleaving when multiple AG2ProfilerHandler instances patch/unpatch; change to use class-level patch state: add a class-level lock (e.g., AG2ProfilerHandler._patch_lock) and a refcount or handler registry (e.g., AG2ProfilerHandler._patch_count or _handlers) and centralize the patch/unpatch logic so _create_llm_wrapper is applied only once on the first patch and the original OpenAIWrapper.create is restored only when the last handler unpatches; update create/close or enable/disable methods to increment/decrement the refcount under the lock and replace instance fields self._original_create/self._instrumented with references into the class-level registry.
🧹 Nitpick comments (1)
docs/source/components/integrations/frameworks.md (1)
94-94: Consider varying sentence openings in this paragraph for readability.Two consecutive sentences begin with “AG2”; a small rephrase would improve flow.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/source/components/integrations/frameworks.md` at line 94, The paragraph repeats the same sentence opening with "AG2" twice; reword the second sentence to vary the opening for better flow—locate the paragraph containing "AG2 is the community-maintained successor..." and change the second sentence that starts with "AG2 provides a modern API..." to something like "It provides a modern API..." or "This framework offers a modern API..." while preserving the content about GroupChat patterns, tool integration, and flexible LLM configuration.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 175-176: The profiler calls like
handler.step_manager.push_intermediate_step(...) and other step_manager
interactions around the AG2 callback wrapper must not be allowed to bubble
exceptions into the core LLM call; wrap each profiler/step_manager invocation
(e.g., push_intermediate_step, finalize/emit steps used in the callback handler)
in a try/except that catches Exception, logs the full traceback via
logger.exception(...), and continues without re-raising so original_func(...) is
never affected by profiler failures; apply this hardening to the occurrences
around the callback wrapper (the push_intermediate_step calls and the blocks
between lines ~182-194 and ~201-216) and ensure the except blocks do not change
control flow or return values of the wrapped original_func call.
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Line 15: Update the docstrings and prose to use repository conventions:
surround code entities like tool_wrapper.py with backticks (e.g.,
`tool_wrapper.py`) and replace uppercase `NAT` with the preferred `nat` or the
full `NeMo Agent Toolkit` wording where applicable (references around the test
module and mentions on Lines ~41 and ~96). Locate occurrences in the test module
(e.g., the module docstring and any inline mentions of NAT or tool_wrapper.py)
and edit them to use backticks for filenames and the canonical `nat`/`NeMo Agent
Toolkit` terms.
---
Duplicate comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py`:
- Around line 81-83: The current monkey-patch in AG2ProfilerHandler stores patch
state per-instance (self._original_create, self._instrumented) and directly
assigns OpenAIWrapper.create, which leads to unsafe interleaving when multiple
AG2ProfilerHandler instances patch/unpatch; change to use class-level patch
state: add a class-level lock (e.g., AG2ProfilerHandler._patch_lock) and a
refcount or handler registry (e.g., AG2ProfilerHandler._patch_count or
_handlers) and centralize the patch/unpatch logic so _create_llm_wrapper is
applied only once on the first patch and the original OpenAIWrapper.create is
restored only when the last handler unpatches; update create/close or
enable/disable methods to increment/decrement the refcount under the lock and
replace instance fields self._original_create/self._instrumented with references
into the class-level registry.
---
Nitpick comments:
In `@docs/source/components/integrations/frameworks.md`:
- Line 94: The paragraph repeats the same sentence opening with "AG2" twice;
reword the second sentence to vary the opening for better flow—locate the
paragraph containing "AG2 is the community-maintained successor..." and change
the second sentence that starts with "AG2 provides a modern API..." to something
like "It provides a modern API..." or "This framework offers a modern API..."
while preserving the content about GroupChat patterns, tool integration, and
flexible LLM configuration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ed31c960-a576-4638-9820-bcb21edd6830
📒 Files selected for processing (4)
docs/source/components/integrations/frameworks.mdexamples/frameworks/nat_ag2_demo/pyproject.tomlpackages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/tests/test_tool_wrapper.py
✅ Files skipped from review due to trivial changes (1)
- examples/frameworks/nat_ag2_demo/pyproject.toml
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.py
Outdated
Show resolved
Hide resolved
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
coderabbitai
bot
added
the
DO NOT MERGE
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py (2)
72-84: Add schema tests for dataclass and Pydantic dataclass inputs.Current coverage only checks
BaseModelandNone. Please add cases for standard dataclasses and Pydantic dataclasses to guard parameter extraction regressions in wrapper schema generation.Based on learnings, parameter extraction must distinguish dataclasses (including Pydantic dataclasses via
dataclasses.fields()) from PydanticBaseModel(model_fields).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py` around lines 72 - 84, Add two tests alongside test_tool_has_json_schema: one that sets mock_function.input_schema to a standard dataclass type and another that sets it to a Pydantic dataclass type, then call ag2_tool_wrapper("my_tool", mock_function, mock_builder) and assert the returned tool has a non-None _func_schema and contains "function"->"parameters"; ensure the tests exercise the branch that distinguishes dataclasses (use dataclasses.fields() detection) from Pydantic BaseModel (model_fields) so parameter extraction logic is covered.
189-207: Path-selection tests are too indirect to verify invoke vs stream routing.Checking only
tool.func.__name__does not prove which underlying method is used. Please assert actual call-path behavior (acall_invokevsacall_stream) after executingtool.func(...).✅ Stronger assertion pattern for non-streaming path
- def test_non_streaming_fn_selects_invoke_path(self, mock_builder): + async def test_non_streaming_fn_selects_invoke_path(self, mock_builder): """Test that non-streaming functions use the invoke path.""" mock_fn = Mock(spec=Function) mock_fn.description = "Non-streaming" mock_fn.input_schema = None mock_fn.has_streaming_output = False mock_fn.has_single_output = True - mock_fn.acall_invoke = AsyncMock() + mock_fn.acall_invoke = AsyncMock(return_value="ok") mock_fn.acall_stream = AsyncMock() tool = ag2_tool_wrapper("invoke_tool", mock_fn, mock_builder) - assert tool.func.__name__ == "invoke_tool" + result = await tool.func(query="q") + assert result == "ok" + mock_fn.acall_invoke.assert_awaited_once_with(query="q") + mock_fn.acall_stream.assert_not_called()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py` around lines 189 - 207, The tests only assert tool.func.__name__ and must be extended to verify actual routing: in test_non_streaming_fn_selects_invoke_path call/await tool.func(...) and assert mock_fn.acall_invoke was awaited (e.g., assert_awaited_once or assert_called_once_with async equivalent) and mock_fn.acall_stream was not awaited; likewise in test_streaming_fn_selects_stream_path call/await tool.func(...) and assert mock_fn.acall_stream was awaited and mock_fn.acall_invoke was not. Use the existing ag2_tool_wrapper, mock_fn, and tool.func symbols to locate where to invoke and assert the correct acall_* methods on the AsyncMock instances.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Around line 156-178: The tests expect streamed string chunks to be
concatenated and an empty stream to return an empty string, but current
ag2_tool_wrapper (in function ag2_tool_wrapper, specifically the code that
collects/returns streaming results used by tool.func) returns a list and None
for empty streams; update the streaming collection logic inside ag2_tool_wrapper
so that if all collected chunks are instances of str it returns "".join(chunks)
and for an empty collection returns "" (not None), while preserving the existing
behavior of returning the list when any chunk is non-string.
---
Nitpick comments:
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Around line 72-84: Add two tests alongside test_tool_has_json_schema: one that
sets mock_function.input_schema to a standard dataclass type and another that
sets it to a Pydantic dataclass type, then call ag2_tool_wrapper("my_tool",
mock_function, mock_builder) and assert the returned tool has a non-None
_func_schema and contains "function"->"parameters"; ensure the tests exercise
the branch that distinguishes dataclasses (use dataclasses.fields() detection)
from Pydantic BaseModel (model_fields) so parameter extraction logic is covered.
- Around line 189-207: The tests only assert tool.func.__name__ and must be
extended to verify actual routing: in test_non_streaming_fn_selects_invoke_path
call/await tool.func(...) and assert mock_fn.acall_invoke was awaited (e.g.,
assert_awaited_once or assert_called_once_with async equivalent) and
mock_fn.acall_stream was not awaited; likewise in
test_streaming_fn_selects_stream_path call/await tool.func(...) and assert
mock_fn.acall_stream was awaited and mock_fn.acall_invoke was not. Use the
existing ag2_tool_wrapper, mock_fn, and tool.func symbols to locate where to
invoke and assert the correct acall_* methods on the AsyncMock instances.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 99cd2399-1fb8-4cfd-b2c9-c9767bbb2e6c
📒 Files selected for processing (2)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/callback_handler.pypackages/nvidia_nat_ag2/tests/test_tool_wrapper.py
Comment on lines
+156
to
+178
| async def test_stream_collected_joins_strings(self, mock_builder): | ||
| """Test that string stream chunks are joined.""" | ||
| mock_fn = self._make_streaming_fn(["Hello", " ", "World"]) | ||
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | ||
|
|
||
| result = await tool.func(query="test") | ||
| assert result == "Hello World" | ||
|
|
||
| async def test_stream_collected_returns_list_for_non_strings(self, mock_builder): | ||
| """Test that non-string stream chunks are returned as list.""" | ||
| mock_fn = self._make_streaming_fn([{"a": 1}, {"b": 2}]) | ||
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | ||
|
|
||
| result = await tool.func(query="test") | ||
| assert result == [{"a": 1}, {"b": 2}] | ||
|
|
||
| async def test_stream_collected_empty_returns_empty_string(self, mock_builder): | ||
| """Test that empty stream returns empty string.""" | ||
| mock_fn = self._make_streaming_fn([]) | ||
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | ||
|
|
||
| result = await tool.func(query="test") | ||
| assert result == "" |
There was a problem hiding this comment.
Streaming assertions conflict with current ag2_tool_wrapper behavior.
These tests expect joined strings and "" for empty streams, but packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py currently returns a collected list and None for empty streams. As written, these tests will fail against the implementation.
🔧 Proposed test fix to match current implementation contract
- async def test_stream_collected_joins_strings(self, mock_builder):
- """Test that string stream chunks are joined."""
+ async def test_stream_collected_returns_list_for_strings(self, mock_builder):
+ """Test that string stream chunks are returned as a list."""
mock_fn = self._make_streaming_fn(["Hello", " ", "World"])
tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder)
result = await tool.func(query="test")
- assert result == "Hello World"
+ assert result == ["Hello", " ", "World"]
@@
- async def test_stream_collected_empty_returns_empty_string(self, mock_builder):
- """Test that empty stream returns empty string."""
+ async def test_stream_collected_empty_returns_none(self, mock_builder):
+ """Test that empty stream returns None."""
mock_fn = self._make_streaming_fn([])
tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder)
result = await tool.func(query="test")
- assert result == ""
+ assert result is None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def test_stream_collected_joins_strings(self, mock_builder): | |
| """Test that string stream chunks are joined.""" | |
| mock_fn = self._make_streaming_fn(["Hello", " ", "World"]) | |
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | |
| result = await tool.func(query="test") | |
| assert result == "Hello World" | |
| async def test_stream_collected_returns_list_for_non_strings(self, mock_builder): | |
| """Test that non-string stream chunks are returned as list.""" | |
| mock_fn = self._make_streaming_fn([{"a": 1}, {"b": 2}]) | |
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | |
| result = await tool.func(query="test") | |
| assert result == [{"a": 1}, {"b": 2}] | |
| async def test_stream_collected_empty_returns_empty_string(self, mock_builder): | |
| """Test that empty stream returns empty string.""" | |
| mock_fn = self._make_streaming_fn([]) | |
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | |
| result = await tool.func(query="test") | |
| assert result == "" | |
| async def test_stream_collected_returns_list_for_strings(self, mock_builder): | |
| """Test that string stream chunks are returned as a list.""" | |
| mock_fn = self._make_streaming_fn(["Hello", " ", "World"]) | |
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | |
| result = await tool.func(query="test") | |
| assert result == ["Hello", " ", "World"] | |
| async def test_stream_collected_returns_list_for_non_strings(self, mock_builder): | |
| """Test that non-string stream chunks are returned as list.""" | |
| mock_fn = self._make_streaming_fn([{"a": 1}, {"b": 2}]) | |
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | |
| result = await tool.func(query="test") | |
| assert result == [{"a": 1}, {"b": 2}] | |
| async def test_stream_collected_empty_returns_none(self, mock_builder): | |
| """Test that empty stream returns None.""" | |
| mock_fn = self._make_streaming_fn([]) | |
| tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder) | |
| result = await tool.func(query="test") | |
| assert result is None |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py` around lines 156 - 178,
The tests expect streamed string chunks to be concatenated and an empty stream
to return an empty string, but current ag2_tool_wrapper (in function
ag2_tool_wrapper, specifically the code that collects/returns streaming results
used by tool.func) returns a list and None for empty streams; update the
streaming collection logic inside ag2_tool_wrapper so that if all collected
chunks are instances of str it returns "".join(chunks) and for an empty
collection returns "" (not None), while preserving the existing behavior of
returning the list when any chunk is non-string.
There was a problem hiding this comment.
if you look closely at the tool wrapper
async def callable_astream_collected(**kwargs: Any) -> Any:
"""Async function to collect all streamed results from the NAT function.
AG2 tools return a single value, so streaming results are collected
and joined into a single response.
Args:
**kwargs: Keyword arguments to pass to the NAT function.
Returns:
Any: The collected streaming results.
"""
chunks: list[Any] = []
async for item in fn.acall_stream(**kwargs):
chunks.append(item)
if not chunks:
return ""
if all(isinstance(c, str) for c in chunks):
return "".join(chunks)
return chunks
It will merge the list into a single string if and only if all the items in the list are strings.
There I will add a test to will handle cross-type situations
There was a problem hiding this comment.
@jsun-m, you're right — I apologize for the incorrect review. Looking at the actual callable_astream_collected implementation you've shared, the existing tests are all correct:
test_stream_collected_joins_strings—"".join(["Hello", " ", "World"])→"Hello World"✅test_stream_collected_returns_list_for_non_strings— non-string chunks fall through toreturn chunks✅test_stream_collected_empty_returns_empty_string— empty list hitsreturn ""✅
I had incorrectly assumed the implementation returned None for empty and a list for all-string chunks without verifying the source. My suggested fix was wrong.
The planned cross-type test (mixed str and non-str chunks) is a great addition — it would verify the "not all strings" branch, e.g.:
async def test_stream_collected_mixed_types_returns_list(self, mock_builder):
"""Test that mixed-type stream chunks are returned as a list."""
mock_fn = self._make_streaming_fn(["text", 42, {"key": "val"}])
tool = ag2_tool_wrapper("stream_tool", mock_fn, mock_builder)
result = await tool.func(query="test")
assert result == ["text", 42, {"key": "val"}]This would cover the edge case where all(isinstance(c, str) for c in chunks) is False due to mixed types.
🧠 Learnings used
Learnt from: AnuradhaKaruppiah
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 1666
File: packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/react_agent/agent.py:377-384
Timestamp: 2026-02-26T17:10:18.070Z
Learning: In `packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/react_agent/agent.py`, the `tool_node` method intentionally strips tool_input via `str(agent_thoughts.tool_input).strip()` and uses the stripped version even when structured parsing fails as fallback—this is pre-existing behavior that should be preserved.
Learnt from: mnajafian-nv
Repo: NVIDIA/NeMo-Agent-Toolkit PR: 1716
File: ci/scripts/model_health_check.py:163-167
Timestamp: 2026-03-03T01:17:03.694Z
Learning: Do not rename external/SDK environment variables when applying the NAT prefix convention. NAT_ prefix is intended only for NAT-internal configuration variables (e.g., NAT_CONFIG_FILE, NAT_LOG_LEVEL, NAT_SPAN_PREFIX) and NAT CI infrastructure variables (e.g., NAT_CI_REDIS_HOST). External SDK environment variables like NVIDIA_API_KEY or OPENAI_API_KEY should remain as-is and not be prefixed with NAT_. When reviewing Python code (including ci/scripts/*.py), ensure SDK-related environment variable names are preserved and not inadvertently renamed to include NAT_. Ensure consistent, explicit usage in code and configuration.
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
coderabbitai
bot
removed
the
DO NOT MERGE
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py (1)
39-49: Consolidate duplicatedmock_functionfixture setup.The two fixture bodies are nearly identical and can be centralized into a shared fixture/helper to reduce drift and maintenance overhead.
As per coding guidelines, “Any frequently repeated code should be extracted into pytest fixtures.”
Also applies to: 94-104
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py` around lines 39 - 49, Extract the duplicated mock setup into a single reusable pytest fixture or helper factory: replace both fixture bodies with a single fixture (e.g., mock_function) or a helper function (e.g., make_mock_function) that constructs the Mock(spec=Function) and sets description, input_schema, has_streaming_output, has_single_output, acall_invoke (AsyncMock returning "test_result") and acall_stream (AsyncMock), then update the other test locations to call the shared fixture/helper; ensure references to mock_function, acall_invoke, acall_stream, description, input_schema, has_streaming_output, and has_single_output remain unchanged so tests keep their behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Around line 199-205: The test test_streaming_fn_selects_stream_path currently
only checks tool.func.__name__ (which both branches set) so update it to assert
runtime branch behavior: create a mock/spy via _make_streaming_fn that records
or returns a distinct marker when invoked in streaming mode, call tool.func (or
run the wrapper) and assert the streaming-marker was produced and the
stream-specific builder path in ag2_tool_wrapper was used; apply the same change
to the companion test in the 206-217 range to verify non-streaming uses the
collect path by asserting its distinct runtime marker/behavior.
---
Nitpick comments:
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Around line 39-49: Extract the duplicated mock setup into a single reusable
pytest fixture or helper factory: replace both fixture bodies with a single
fixture (e.g., mock_function) or a helper function (e.g., make_mock_function)
that constructs the Mock(spec=Function) and sets description, input_schema,
has_streaming_output, has_single_output, acall_invoke (AsyncMock returning
"test_result") and acall_stream (AsyncMock), then update the other test
locations to call the shared fixture/helper; ensure references to mock_function,
acall_invoke, acall_stream, description, input_schema, has_streaming_output, and
has_single_output remain unchanged so tests keep their behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f8d5597e-f734-4bca-b4b6-9ff4fca9e604
📒 Files selected for processing (1)
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
coderabbitai
bot
added
the
DO NOT MERGE
There was a problem hiding this comment.
🧹 Nitpick comments (1)
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py (1)
29-34: Add schema-path tests for dataclass variants to prevent wrapper regressions.
input_schemacoverage is currently limited toBaseModelandNone. Please also add tests for stdlib dataclass and Pydantic dataclass schemas, since parameter extraction logic often differs across those types.Proposed test additions
+from dataclasses import dataclass + import pytest from pydantic import BaseModel +from pydantic.dataclasses import dataclass as pydantic_dataclass @@ class MockInputSchema(BaseModel): @@ query: str limit: int = 10 + + +@dataclass +class MockDataclassInputSchema: + """Mock stdlib dataclass input schema for tool wrapper.""" + + query: str + limit: int = 10 + + +@pydantic_dataclass +class MockPydanticDataclassInputSchema: + """Mock pydantic dataclass input schema for tool wrapper.""" + + query: str + limit: int = 10 @@ def test_tool_has_json_schema(self, mock_function, mock_builder): """Test that created Tool has parameter schema from input_schema.""" tool = ag2_tool_wrapper("my_tool", mock_function, mock_builder) schema = tool._func_schema assert schema is not None assert schema["function"]["parameters"] is not None + + def test_tool_has_json_schema_for_stdlib_dataclass(self, mock_function, mock_builder): + """Test that schema extraction works for stdlib dataclass input schemas.""" + mock_function.input_schema = MockDataclassInputSchema + tool = ag2_tool_wrapper("my_tool", mock_function, mock_builder) + assert tool._func_schema is not None + assert tool._func_schema["function"]["parameters"] is not None + + def test_tool_has_json_schema_for_pydantic_dataclass(self, mock_function, mock_builder): + """Test that schema extraction works for pydantic dataclass input schemas.""" + mock_function.input_schema = MockPydanticDataclassInputSchema + tool = ag2_tool_wrapper("my_tool", mock_function, mock_builder) + assert tool._func_schema is not None + assert tool._func_schema["function"]["parameters"] is not NoneBased on learnings, in
packages/nvidia_nat_autogen/src/nat/plugins/autogen/tool_wrapper.pyparameter extraction had to treat dataclasses andBaseModeldifferently (dataclasses.fields()vsmodel_fields/__pydantic_fields__).Also applies to: 72-83
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py` around lines 29 - 34, Add tests in packages/nvidia_nat_ag2/tests/test_tool_wrapper.py that cover dataclass variants of input_schema in addition to MockInputSchema (Pydantic BaseModel) — specifically create a stdlib `@dataclass` (with fields like query, limit) and a Pydantic dataclass (using pydantic.dataclasses.dataclass) and assert the tool wrapper extracts parameters the same way as for MockInputSchema; this ensures the extraction logic in packages/nvidia_nat_autogen/src/nat/plugins/autogen/tool_wrapper.py handles dataclasses (use dataclasses.fields()) and Pydantic models (use model_fields or __pydantic_fields__) correctly and prevents regressions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Around line 29-34: Add tests in
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py that cover dataclass variants
of input_schema in addition to MockInputSchema (Pydantic BaseModel) —
specifically create a stdlib `@dataclass` (with fields like query, limit) and a
Pydantic dataclass (using pydantic.dataclasses.dataclass) and assert the tool
wrapper extracts parameters the same way as for MockInputSchema; this ensures
the extraction logic in
packages/nvidia_nat_autogen/src/nat/plugins/autogen/tool_wrapper.py handles
dataclasses (use dataclasses.fields()) and Pydantic models (use model_fields or
__pydantic_fields__) correctly and prevents regressions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f66247a6-5ede-4bad-b883-950d08bc19e9
📒 Files selected for processing (1)
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py
|
@coderabbitai resume |
✅ Actions performedReviews resumed. |
Signed-off-by: jsun-m <91754185+jsun-m@users.noreply.github.com>
coderabbitai
bot
removed
the
DO NOT MERGE
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
packages/nvidia_nat_ag2/tests/test_tool_wrapper.py (1)
59-73: Add type hints to fixtures/helpers to match repository typing policy.Several reusable fixture/helper methods (for example Lines 59, 71, 128, 140, 170, 173) are unannotated. Adding explicit parameter/return typing will align better with the project’s Python typing requirements.
Proposed patch
+from collections.abc import Sequence +from typing import Any @@ - def fixture_mock_function(self): + def fixture_mock_function(self) -> Mock: @@ - def fixture_mock_builder(self): + def fixture_mock_builder(self) -> Mock: @@ - def fixture_mock_function(self): + def fixture_mock_function(self) -> Mock: @@ - def fixture_mock_builder(self): + def fixture_mock_builder(self) -> Mock: @@ - def fixture_mock_builder(self): + def fixture_mock_builder(self) -> Mock: @@ - def _make_streaming_fn(self, stream_items): + def _make_streaming_fn(self, stream_items: Sequence[Any]) -> Mock:As per coding guidelines: "Python methods should use type hints for all parameters and return values" and "All public APIs require Python 3.11+ type hints on parameters and return values".
Also applies to: 127-141, 169-174
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py` around lines 59 - 73, Add explicit type hints to the unannotated fixture/helper functions: annotate fixture_mock_function to return Mock (or Function) with proper typing (e.g., -> Mock[Function] or -> Function) and any parameters (none) as -> Function/Mock; annotate fixture_mock_builder to return Mock[Builder] or Builder (-> Mock[Builder]); ensure MockInputSchema/AsyncMock usages have compatible return type annotations for acall_invoke (e.g., AsyncMock returning str) if you add local helper stubs; update imports if needed to reference typing.Generic/typing.Any/typing.TypeAlias. Locate and update the functions named fixture_mock_function and fixture_mock_builder (and the other listed helpers at lines ~127-141 and ~169-174) to include full parameter and return type hints consistent with the repository policy.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.py`:
- Around line 15-73: Update wording and docstrings to use the repo conventions:
replace capitalized "NAT" with lowercase "nat" in the module docstring and in
all public-facing docstrings, and convert the ag2_tool_wrapper public docstring
to a Google-style docstring that describes Args and Returns; also update the
inner docstrings for callable_ainvoke and callable_astream_collected to match
repository wording (use "nat" and brief Google-style parameter/return
descriptions). Target the module-level string at the top and the docstrings for
ag2_tool_wrapper, callable_ainvoke, callable_astream_collected (and any other
occurrences found by searching for "NAT") so naming and style conform to repo
guidelines.
---
Nitpick comments:
In `@packages/nvidia_nat_ag2/tests/test_tool_wrapper.py`:
- Around line 59-73: Add explicit type hints to the unannotated fixture/helper
functions: annotate fixture_mock_function to return Mock (or Function) with
proper typing (e.g., -> Mock[Function] or -> Function) and any parameters (none)
as -> Function/Mock; annotate fixture_mock_builder to return Mock[Builder] or
Builder (-> Mock[Builder]); ensure MockInputSchema/AsyncMock usages have
compatible return type annotations for acall_invoke (e.g., AsyncMock returning
str) if you add local helper stubs; update imports if needed to reference
typing.Generic/typing.Any/typing.TypeAlias. Locate and update the functions
named fixture_mock_function and fixture_mock_builder (and the other listed
helpers at lines ~127-141 and ~169-174) to include full parameter and return
type hints consistent with the repository policy.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 1ba1c8f9-d0cf-4fc4-b7e1-a4c5791235c5
📒 Files selected for processing (2)
packages/nvidia_nat_ag2/src/nat/plugins/ag2/tool_wrapper.pypackages/nvidia_nat_ag2/tests/test_tool_wrapper.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds AG2 as a supported framework in NeMo-Agent-Toolkit alongside the existing Microsoft AutoGen integration. AG2 is the community- driven fork of Microsoft AutoGen with a modernized API for multi- agent orchestration.
The two packages coexist — they target different APIs (AG2 0.11+ vs Microsoft AutoGen 0.7.x) and different package names (ag2 vs autogen-agentchat).
Changes:
tool wrapper (NAT Function → AG2 Tool), profiler callback
NAT profiling integration
Addresses community demand from Issue #366.
Summary by CodeRabbit