From bcf24530d9639ca0142679cf9651fcc7052ddafb Mon Sep 17 00:00:00 2001 From: Daniel Wang Date: Thu, 25 Sep 2025 15:52:10 -0700 Subject: [PATCH 1/6] initial code Signed-off-by: Daniel Wang --- .../commands/workflow/templates/config.yml.j2 | 7 +++-- .../workflow/templates/workflow.py.j2 | 28 ++++++------------- 2 files changed, 13 insertions(+), 22 deletions(-) diff --git a/src/nat/cli/commands/workflow/templates/config.yml.j2 b/src/nat/cli/commands/workflow/templates/config.yml.j2 index 4e57fd2ea..39c7cb96c 100644 --- a/src/nat/cli/commands/workflow/templates/config.yml.j2 +++ b/src/nat/cli/commands/workflow/templates/config.yml.j2 @@ -11,5 +11,8 @@ general: _type: console workflow: - _type: {{workflow_name}} - parameter: default_value + _type: react_agent + workflow_alias: {{workflow_name}} + description: "ReAct-based Workflow" + llm_name: nim_llm + tool_names: [] diff --git a/src/nat/cli/commands/workflow/templates/workflow.py.j2 b/src/nat/cli/commands/workflow/templates/workflow.py.j2 index c48761885..95b9113b7 100644 --- a/src/nat/cli/commands/workflow/templates/workflow.py.j2 +++ b/src/nat/cli/commands/workflow/templates/workflow.py.j2 @@ -1,36 +1,24 @@ import logging -from pydantic import Field - from nat.builder.builder import Builder -from nat.builder.function_info import FunctionInfo +from nat.builder.framework_enum import LLMFrameworkEnum from nat.cli.register_workflow import register_function -from nat.data_models.function import FunctionBaseConfig +from nat.agent.react_agent.register import ReActAgentWorkflowConfig, react_agent_workflow logger = logging.getLogger(__name__) -class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"): +class {{ workflow_class_name }}(ReActAgentWorkflowConfig, name="{{ workflow_name }}"): """ {{workflow_description}} """ - # Add your custom configuration parameters here - parameter: str = Field(default="default_value", description="Notional description for this parameter") + # You can add or override ReAct config parameters here if needed -@register_function(config_type={{ workflow_class_name }}) +@register_function(config_type={{ workflow_class_name }}, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) async def {{ python_safe_workflow_name }}_function( config: {{ workflow_class_name }}, builder: Builder ): - # Implement your function logic here - async def _response_fn(input_message: str) -> str: - # Process the input_message and generate output - output_message = f"Hello from {{ workflow_name }} workflow! You said: {input_message}" - return output_message - - try: - yield FunctionInfo.create(single_fn=_response_fn) - except GeneratorExit: - logger.warning("Function exited early!") - finally: - logger.info("Cleaning up {{ workflow_name }} workflow.") + # Delegate to the built-in ReAct agent workflow + async for info in react_agent_workflow(config, builder): + yield info From d333c32fa23807c16024b9097942a019ea8a9ae8 Mon Sep 17 00:00:00 2001 From: Daniel Wang Date: Thu, 25 Sep 2025 16:12:05 -0700 Subject: [PATCH 2/6] More fixes Signed-off-by: Daniel Wang --- .../cli/commands/workflow/templates/config.yml.j2 | 12 +++++++++++- .../cli/commands/workflow/templates/register.py.j2 | 4 ++-- src/nat/cli/commands/workflow/workflow_commands.py | 2 +- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/nat/cli/commands/workflow/templates/config.yml.j2 b/src/nat/cli/commands/workflow/templates/config.yml.j2 index 39c7cb96c..779369c36 100644 --- a/src/nat/cli/commands/workflow/templates/config.yml.j2 +++ b/src/nat/cli/commands/workflow/templates/config.yml.j2 @@ -10,9 +10,19 @@ general: front_end: _type: console +functions: + current_datetime: + _type: current_datetime + +llms: + nim_llm: + _type: nim + model_name: meta/llama-3.1-8b-instruct + temperature: 0.0 + workflow: _type: react_agent workflow_alias: {{workflow_name}} description: "ReAct-based Workflow" llm_name: nim_llm - tool_names: [] + tool_names: [current_datetime] diff --git a/src/nat/cli/commands/workflow/templates/register.py.j2 b/src/nat/cli/commands/workflow/templates/register.py.j2 index 2b0c8a2a9..44b66e3a0 100644 --- a/src/nat/cli/commands/workflow/templates/register.py.j2 +++ b/src/nat/cli/commands/workflow/templates/register.py.j2 @@ -1,4 +1,4 @@ # flake8: noqa -# Import any tools which need to be automatically registered here -from {{package_name}} import {{workflow_name}}_function +# Import the generated workflow function to trigger registration +from {{ python_safe_workflow_name }} import {{ python_safe_workflow_name }}_function diff --git a/src/nat/cli/commands/workflow/workflow_commands.py b/src/nat/cli/commands/workflow/workflow_commands.py index 7abf3522e..99edaac6a 100644 --- a/src/nat/cli/commands/workflow/workflow_commands.py +++ b/src/nat/cli/commands/workflow/workflow_commands.py @@ -224,7 +224,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip files_to_render = { 'pyproject.toml.j2': new_workflow_dir / 'pyproject.toml', 'register.py.j2': base_dir / 'register.py', - 'workflow.py.j2': base_dir / f'{workflow_name}_function.py', + 'workflow.py.j2': base_dir / f'{package_name}_function.py', '__init__.py.j2': base_dir / '__init__.py', 'config.yml.j2': configs_dir / 'config.yml', } From 01ce215a52995ded1ab444505a109b4c2428e363 Mon Sep 17 00:00:00 2001 From: Yuchen Zhang <134643420+yczhang-nv@users.noreply.github.com> Date: Thu, 25 Sep 2025 16:08:38 -0700 Subject: [PATCH 3/6] Skip some tests in `test_mcp_client_base.py` to avoid blocking CI (#850) Some tests with `asyncio.sleep()` is randomly failing in CI tasks. Skip those to avoid them blocking CI. ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit * **Tests** * Temporarily skipped two reconnect backoff timing tests to reduce CI flakiness in continuous integration. * Improves reliability of pipeline runs while investigation continues. * No impact on application functionality or user experience; all other tests remain unchanged and continue to validate core behavior. Authors: - Yuchen Zhang (https://github.com/yczhang-nv) Approvers: - Will Killian (https://github.com/willkill07) URL: https://github.com/NVIDIA/NeMo-Agent-Toolkit/pull/850 Signed-off-by: Zhongxuan (Daniel) Wang --- tests/nat/mcp/test_mcp_client_base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/nat/mcp/test_mcp_client_base.py b/tests/nat/mcp/test_mcp_client_base.py index cbb47d42f..b80f08de2 100644 --- a/tests/nat/mcp/test_mcp_client_base.py +++ b/tests/nat/mcp/test_mcp_client_base.py @@ -315,6 +315,7 @@ async def always_fail(): await client.get_tools() +@pytest.mark.skip(reason="This test might fail in CI due to race conditions") async def test_reconnect_backoff_timing(): """Test that reconnect backoff timing works correctly.""" client = MockMCPClient(transport="streamable-http", @@ -358,6 +359,7 @@ async def mock_list_tools(): assert attempt_times[1] == 0.2 +@pytest.mark.skip(reason="This test might fail in CI due to race conditions") async def test_reconnect_max_backoff_limit(): """Test that backoff doesn't exceed maximum.""" client = MockMCPClient( From 7fd82cd55f7da790ab30707e07e948e90d59632e Mon Sep 17 00:00:00 2001 From: "Zhongxuan (Daniel) Wang" Date: Fri, 26 Sep 2025 21:30:41 +0000 Subject: [PATCH 4/6] Added better templates Signed-off-by: Daniel Wang Signed-off-by: Zhongxuan (Daniel) Wang --- .../commands/workflow/templates/config.yml.j2 | 10 ++++--- .../workflow/templates/register.py.j2 | 2 +- .../workflow/templates/workflow.py.j2 | 26 +++++++++---------- 3 files changed, 20 insertions(+), 18 deletions(-) diff --git a/src/nat/cli/commands/workflow/templates/config.yml.j2 b/src/nat/cli/commands/workflow/templates/config.yml.j2 index 779369c36..d8df4a4f4 100644 --- a/src/nat/cli/commands/workflow/templates/config.yml.j2 +++ b/src/nat/cli/commands/workflow/templates/config.yml.j2 @@ -13,16 +13,18 @@ general: functions: current_datetime: _type: current_datetime + {{python_safe_workflow_name}}_echo: + _type: {{python_safe_workflow_name}}_echo + prefix: "Hello:" llms: nim_llm: _type: nim - model_name: meta/llama-3.1-8b-instruct + model_name: meta/llama-3.1-70b-instruct temperature: 0.0 workflow: _type: react_agent - workflow_alias: {{workflow_name}} - description: "ReAct-based Workflow" llm_name: nim_llm - tool_names: [current_datetime] + tool_names: [current_datetime, {{python_safe_workflow_name}}_echo] + parse_agent_response_max_retries: 3 diff --git a/src/nat/cli/commands/workflow/templates/register.py.j2 b/src/nat/cli/commands/workflow/templates/register.py.j2 index 44b66e3a0..5d44d0acf 100644 --- a/src/nat/cli/commands/workflow/templates/register.py.j2 +++ b/src/nat/cli/commands/workflow/templates/register.py.j2 @@ -1,4 +1,4 @@ # flake8: noqa # Import the generated workflow function to trigger registration -from {{ python_safe_workflow_name }} import {{ python_safe_workflow_name }}_function +from {{package_name}} import {{ python_safe_workflow_name }}_function diff --git a/src/nat/cli/commands/workflow/templates/workflow.py.j2 b/src/nat/cli/commands/workflow/templates/workflow.py.j2 index 95b9113b7..6187e1022 100644 --- a/src/nat/cli/commands/workflow/templates/workflow.py.j2 +++ b/src/nat/cli/commands/workflow/templates/workflow.py.j2 @@ -1,24 +1,24 @@ import logging from nat.builder.builder import Builder +from nat.builder.function_info import FunctionInfo from nat.builder.framework_enum import LLMFrameworkEnum from nat.cli.register_workflow import register_function -from nat.agent.react_agent.register import ReActAgentWorkflowConfig, react_agent_workflow +from nat.data_models.function import FunctionBaseConfig +from pydantic import Field logger = logging.getLogger(__name__) +class {{ workflow_class_name }}EchoConfig(FunctionBaseConfig, name="{{ python_safe_workflow_name }}_echo"): + prefix: str = Field(default="Echo:", description="Prefix to add before the echoed text.") -class {{ workflow_class_name }}(ReActAgentWorkflowConfig, name="{{ workflow_name }}"): - """ - {{workflow_description}} - """ - # You can add or override ReAct config parameters here if needed +@register_function(config_type={{ workflow_class_name }}EchoConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) +async def register_{{ python_safe_workflow_name }}_echo(config: {{ workflow_class_name }}EchoConfig, builder: Builder): -@register_function(config_type={{ workflow_class_name }}, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) -async def {{ python_safe_workflow_name }}_function( - config: {{ workflow_class_name }}, builder: Builder -): - # Delegate to the built-in ReAct agent workflow - async for info in react_agent_workflow(config, builder): - yield info + async def _echo(text: str) -> str: + return f"{config.prefix} {text}" + + yield FunctionInfo.from_fn( + _echo, + description="This is a function that takes an input and echos back with a pre-defined prefix.") From 2904508df8800db9868219db2fc5723de3c7bd39 Mon Sep 17 00:00:00 2001 From: "Zhongxuan (Daniel) Wang" Date: Tue, 30 Sep 2025 17:09:28 -0700 Subject: [PATCH 5/6] Update src/nat/cli/commands/workflow/workflow_commands.py Co-authored-by: Will Killian <2007799+willkill07@users.noreply.github.com> Signed-off-by: Zhongxuan (Daniel) Wang --- src/nat/cli/commands/workflow/workflow_commands.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nat/cli/commands/workflow/workflow_commands.py b/src/nat/cli/commands/workflow/workflow_commands.py index 99edaac6a..1e3359e0e 100644 --- a/src/nat/cli/commands/workflow/workflow_commands.py +++ b/src/nat/cli/commands/workflow/workflow_commands.py @@ -224,7 +224,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip files_to_render = { 'pyproject.toml.j2': new_workflow_dir / 'pyproject.toml', 'register.py.j2': base_dir / 'register.py', - 'workflow.py.j2': base_dir / f'{package_name}_function.py', + 'workflow.py.j2': base_dir / f'{python_safe_workflow_name}_function.py', '__init__.py.j2': base_dir / '__init__.py', 'config.yml.j2': configs_dir / 'config.yml', } From 7222ed3a7958a26bdc4fea00a1573eece0292e2f Mon Sep 17 00:00:00 2001 From: Daniel Wang Date: Wed, 1 Oct 2025 16:31:57 -0700 Subject: [PATCH 6/6] Fixing the feedbacks and some other nitty fixies :) Signed-off-by: Daniel Wang --- .../commands/workflow/templates/config.yml.j2 | 6 ++--- .../workflow/templates/register.py.j2 | 2 +- .../workflow/templates/workflow.py.j2 | 22 +++++++++++++------ .../commands/workflow/workflow_commands.py | 6 +++-- 4 files changed, 23 insertions(+), 13 deletions(-) diff --git a/src/nat/cli/commands/workflow/templates/config.yml.j2 b/src/nat/cli/commands/workflow/templates/config.yml.j2 index d8df4a4f4..918e53077 100644 --- a/src/nat/cli/commands/workflow/templates/config.yml.j2 +++ b/src/nat/cli/commands/workflow/templates/config.yml.j2 @@ -13,8 +13,8 @@ general: functions: current_datetime: _type: current_datetime - {{python_safe_workflow_name}}_echo: - _type: {{python_safe_workflow_name}}_echo + {{python_safe_workflow_name}}: + _type: {{python_safe_workflow_name}} prefix: "Hello:" llms: @@ -26,5 +26,5 @@ llms: workflow: _type: react_agent llm_name: nim_llm - tool_names: [current_datetime, {{python_safe_workflow_name}}_echo] + tool_names: [current_datetime, {{python_safe_workflow_name}}] parse_agent_response_max_retries: 3 diff --git a/src/nat/cli/commands/workflow/templates/register.py.j2 b/src/nat/cli/commands/workflow/templates/register.py.j2 index 5d44d0acf..d5085beef 100644 --- a/src/nat/cli/commands/workflow/templates/register.py.j2 +++ b/src/nat/cli/commands/workflow/templates/register.py.j2 @@ -1,4 +1,4 @@ # flake8: noqa # Import the generated workflow function to trigger registration -from {{package_name}} import {{ python_safe_workflow_name }}_function +from {{package_name}} import {{ python_safe_workflow_name }} diff --git a/src/nat/cli/commands/workflow/templates/workflow.py.j2 b/src/nat/cli/commands/workflow/templates/workflow.py.j2 index 6187e1022..9f07e049d 100644 --- a/src/nat/cli/commands/workflow/templates/workflow.py.j2 +++ b/src/nat/cli/commands/workflow/templates/workflow.py.j2 @@ -1,24 +1,32 @@ import logging +from pydantic import Field + from nat.builder.builder import Builder -from nat.builder.function_info import FunctionInfo from nat.builder.framework_enum import LLMFrameworkEnum +from nat.builder.function_info import FunctionInfo from nat.cli.register_workflow import register_function from nat.data_models.function import FunctionBaseConfig -from pydantic import Field logger = logging.getLogger(__name__) -class {{ workflow_class_name }}EchoConfig(FunctionBaseConfig, name="{{ python_safe_workflow_name }}_echo"): + +class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"): + """ + {{ workflow_description }} + """ prefix: str = Field(default="Echo:", description="Prefix to add before the echoed text.") -@register_function(config_type={{ workflow_class_name }}EchoConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) -async def register_{{ python_safe_workflow_name }}_echo(config: {{ workflow_class_name }}EchoConfig, builder: Builder): +@register_function(config_type={{ workflow_class_name }}, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) +async def {{ python_safe_workflow_name }}_function(config: {{ workflow_class_name }}, builder: Builder): async def _echo(text: str) -> str: + """ + This is the part where all the logic for the function goes. + """ return f"{config.prefix} {text}" + # This is the part where the function is registered. The description parameter is used to describe the function. yield FunctionInfo.from_fn( - _echo, - description="This is a function that takes an input and echos back with a pre-defined prefix.") + _echo, description="This is a function that takes an input and echos back with a pre-defined prefix.") diff --git a/src/nat/cli/commands/workflow/workflow_commands.py b/src/nat/cli/commands/workflow/workflow_commands.py index 1e3359e0e..71edd9a53 100644 --- a/src/nat/cli/commands/workflow/workflow_commands.py +++ b/src/nat/cli/commands/workflow/workflow_commands.py @@ -220,11 +220,13 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip else: install_cmd = ['pip', 'install', '-e', str(new_workflow_dir)] + python_safe_workflow_name = workflow_name.replace("-", "_") + # List of templates and their destinations files_to_render = { 'pyproject.toml.j2': new_workflow_dir / 'pyproject.toml', 'register.py.j2': base_dir / 'register.py', - 'workflow.py.j2': base_dir / f'{python_safe_workflow_name}_function.py', + 'workflow.py.j2': base_dir / f'{python_safe_workflow_name}.py', '__init__.py.j2': base_dir / '__init__.py', 'config.yml.j2': configs_dir / 'config.yml', } @@ -233,7 +235,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip context = { 'editable': editable, 'workflow_name': workflow_name, - 'python_safe_workflow_name': workflow_name.replace("-", "_"), + 'python_safe_workflow_name': python_safe_workflow_name, 'package_name': package_name, 'rel_path_to_repo_root': rel_path_to_repo_root, 'workflow_class_name': f"{_generate_valid_classname(workflow_name)}FunctionConfig",