Skip to content

Commit 94901ef

Browse files
authored
fix(builder): let MCP servers publish group functions without include (#2456)
`nat mcp serve` and `nat fastmcp server run` start with a function group that has no `include`, and publish the functions it makes available as tools. Both servers already list every function a group makes available as a tool, and build one workflow per tool with `SessionManager.create(entry_function=name)`. `WorkflowBuilder.build` resolved `entry_function` only through `get_function`, the global registry, which holds just the `include`d functions, so any other group function failed with "Function `group__fn` not found" and the server didn't start. - `WorkflowBuilder.build` and `PerUserWorkflowBuilder.build` now also resolve an entry function that a function group makes available as a tool (`get_accessible_functions`), when it isn't in the registry. The per-user builder looks in its per-user groups, then the shared ones. - `get_function` is unchanged, so a group function without `include` still isn't globally addressable, as the function group docs describe. An excluded function, or a name no group has, still raises the same "not found" error. If two groups offer the same full name (group `a` with `b__c`, group `a__b` with `c`), the build fails with an "ambiguous" error instead of picking one. - This applies to every `build(entry_function=...)` caller, so a FastAPI endpoint's `function_name` can also name such a function. The function group docs now say so under the include/exclude table. - Legacy dotted names (`group.fn`) aren't resolved for these functions; the servers always pass `__` names, and dotted names are slated for removal (TODO #1952). Tests (fail on develop, except the `include` cases, which pass on both): - `nvidia_nat_core/tests/nat/builder/test_function_group_entry_function.py`: through `WorkflowBuilder` and `SessionManager.create(entry_function=...)`, a group function is the entry point with no filter, with `exclude` and with `include`, also for a per-user workflow with a shared or a per-user group. Excluded and unknown names still aren't found, and two groups offering the same name is an error. - `nvidia_nat_mcp/tests/server/test_mcp_publishes_group_functions.py`: the MCP worker with a real builder and `FastMCP` publishes the group's tools with `tool_names` empty, a group name, a function name, `exclude` and `include`, and calling `text__upper` returns the result. `nvidia_nat_fastmcp/tests/server/test_fastmcp_publishes_group_functions.py` does the same for the FastMCP worker through an in-memory `fastmcp.Client`. End to end with the core `github` group (no `include`) and an MCP client listing tools: on develop both `nat mcp serve` and `nat fastmcp server run` exit with "Function `github__...` not found"; on this branch both start and list all 6 `github__*` tools. `nvidia_nat_core` builder and runtime tests, `nvidia_nat_mcp/tests/server` and `nvidia_nat_fastmcp/tests` pass. yapf and ruff pass. Closes #2455 ## 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/index.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit * **New Features** * Functions exposed through function groups can now be used as workflow entry points, including in per-user workflows. * Function-group tools are available through MCP and FastMCP, with support for selecting or filtering which tools are published. * Ambiguous function names across groups are reported as errors. * **Documentation** * Clarified that group-exposed functions can serve as workflow entry points without being registered globally. Authors: - Somay kousis (https://github.com/Somay-kousis) Approvers: - Will Killian (https://github.com/willkill07) URL: #2456
1 parent c57e919 commit 94901ef

6 files changed

Lines changed: 411 additions & 3 deletions

File tree

‎docs/source/build-workflows/functions-and-function-groups/function-groups.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -458,6 +458,8 @@ function_groups:
458458
| `include: [add]` | ✓ (all functions) | ✓ (only `add`) | ✓ (only `add`) |
459459
| `exclude: [divide]` | ✓ (all functions) | ✗ | ✓ (except `divide`) |
460460

461+
A function that is available as an agent tool can also be the entry point of a workflow, even when it is not in the global registry. For example, `nat mcp serve` publishes each of them as an MCP tool, and a FastAPI endpoint can name one in `function_name`.
462+
461463
## Using Function Groups
462464

463465
### Creating Custom Function Groups

‎packages/nvidia_nat_core/src/nat/builder/per_user_workflow_builder.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
from nat.builder.workflow_builder import WorkflowBuilder
3838
from nat.builder.workflow_builder import _build_function_group_impl
3939
from nat.builder.workflow_builder import _build_function_impl
40+
from nat.builder.workflow_builder import _find_function_group_function
4041
from nat.builder.workflow_builder import _log_build_failure
4142
from nat.cli.type_registry import GlobalTypeRegistry
4243
from nat.cli.type_registry import TypeRegistry
@@ -697,8 +698,15 @@ async def build(self, entry_function: str | None = None) -> Workflow:
697698
# Use workflow (could be per-user or shared)
698699
entry_fn_obj = self.get_workflow()
699700
else:
700-
# Use specified function (could be per-user or shared)
701-
entry_fn_obj = await self.get_function(entry_function)
701+
# Use specified function (could be per-user or shared, or a function a group makes available as a tool)
702+
try:
703+
entry_fn_obj = await self.get_function(entry_function)
704+
except ValueError:
705+
entry_fn_obj = await _find_function_group_function(
706+
entry_function,
707+
[*self._per_user_function_groups.values(), *self._shared_builder._function_groups.values()])
708+
if entry_fn_obj is None:
709+
raise
702710

703711
# Collect function names that are included by function groups (shared + per-user)
704712
# These will be skipped when populating function_configs and all_functions

‎packages/nvidia_nat_core/src/nat/builder/workflow_builder.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,19 @@ class ConfiguredTrajectoryBuilder:
189189
instance: TrajectoryBuilder
190190

191191

192+
async def _find_function_group_function(name: str, function_groups: list[ConfiguredFunctionGroup]) -> Function | None:
193+
"""A function one of the groups makes available as a tool (`group__fn`), even if it isn't globally registered."""
194+
matches = []
195+
for function_group in function_groups:
196+
accessible_functions = await function_group.instance.get_accessible_functions()
197+
if name in accessible_functions:
198+
matches.append(accessible_functions[name])
199+
if len(matches) > 1:
200+
# e.g. group `a` with function `b__c` and group `a__b` with function `c`
201+
raise ValueError(f"Function `{name}` is ambiguous: more than one function group makes it available")
202+
return matches[0] if matches else None
203+
204+
192205
def _log_build_failure(component_name: str,
193206
component_type: str,
194207
completed_components: list[tuple[str, str]],
@@ -538,7 +551,7 @@ async def build(self, entry_function: str | None = None) -> Workflow:
538551
if (entry_function is None):
539552
entry_fn_obj = self.get_workflow()
540553
else:
541-
entry_fn_obj = await self.get_function(entry_function)
554+
entry_fn_obj = await self._get_entry_function(entry_function)
542555

543556
workflow = Workflow.from_entry_fn(config=config,
544557
entry_fn=entry_fn_obj,
@@ -703,6 +716,16 @@ async def add_function_group(self, name: str | FunctionGroupRef, config: Functio
703716

704717
return build_result.instance
705718

719+
async def _get_entry_function(self, name: str) -> Function:
720+
"""A registered function, or one a function group makes available as a tool (MCP servers publish both)."""
721+
try:
722+
return await self.get_function(name)
723+
except ValueError:
724+
group_function = await _find_function_group_function(name, list(self._function_groups.values()))
725+
if group_function is None:
726+
raise
727+
return group_function
728+
706729
def _check_backwards_compatibility_function_name(self, name: str) -> str:
707730
if name in self._functions:
708731
return name
Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
"""A function a group makes available as a tool can be the entry point of a workflow (as MCP servers do)."""
16+
17+
import pytest
18+
19+
from nat.builder.builder import Builder
20+
from nat.builder.function import FunctionGroup
21+
from nat.builder.function_info import FunctionInfo
22+
from nat.builder.workflow_builder import WorkflowBuilder
23+
from nat.cli.register_workflow import register_function
24+
from nat.cli.register_workflow import register_function_group
25+
from nat.cli.register_workflow import register_per_user_function
26+
from nat.cli.register_workflow import register_per_user_function_group
27+
from nat.data_models.config import Config
28+
from nat.data_models.function import FunctionBaseConfig
29+
from nat.data_models.function import FunctionGroupBaseConfig
30+
from nat.runtime.session import SessionManager
31+
32+
33+
class TextToolsConfig(FunctionGroupBaseConfig, name="entry_test_text_tools"):
34+
pass
35+
36+
37+
class EchoWorkflowConfig(FunctionBaseConfig, name="entry_test_echo_workflow"):
38+
pass
39+
40+
41+
class PerUserTextToolsConfig(FunctionGroupBaseConfig, name="entry_test_per_user_text_tools"):
42+
pass
43+
44+
45+
class PerUserEchoWorkflowConfig(FunctionBaseConfig, name="entry_test_per_user_echo_workflow"):
46+
pass
47+
48+
49+
def _text_tools_group(config) -> FunctionGroup:
50+
group = FunctionGroup(config=config)
51+
52+
async def upper(text: str) -> str:
53+
return text.upper()
54+
55+
async def reverse(text: str) -> str:
56+
return text[::-1]
57+
58+
group.add_function("upper", upper, description="upper-case the text")
59+
group.add_function("reverse", reverse, description="reverse the text")
60+
return group
61+
62+
63+
@pytest.fixture(name="register_components", scope="module", autouse=True)
64+
def register_components_fixture():
65+
66+
@register_function_group(config_type=TextToolsConfig)
67+
async def _text_tools(config: TextToolsConfig, builder: Builder):
68+
yield _text_tools_group(config)
69+
70+
@register_per_user_function_group(config_type=PerUserTextToolsConfig)
71+
async def _per_user_text_tools(config: PerUserTextToolsConfig, builder: Builder):
72+
yield _text_tools_group(config)
73+
74+
@register_per_user_function(config_type=PerUserEchoWorkflowConfig, input_type=str, single_output_type=str)
75+
async def _per_user_workflow(config, builder: Builder):
76+
77+
async def _inner(text: str) -> str:
78+
return text
79+
80+
yield FunctionInfo.from_fn(_inner, description="echo")
81+
82+
@register_function(config_type=EchoWorkflowConfig)
83+
async def _workflow(config, builder: Builder):
84+
85+
async def _inner(text: str) -> str:
86+
return text
87+
88+
yield FunctionInfo.from_fn(_inner, description="echo")
89+
90+
91+
def _config(group_type: str = "entry_test_text_tools",
92+
workflow_type: str = "entry_test_echo_workflow",
93+
**group_options) -> Config:
94+
return Config.model_validate({
95+
"function_groups": {
96+
"text": {
97+
"_type": group_type, **group_options
98+
}
99+
},
100+
"workflow": {
101+
"_type": workflow_type
102+
},
103+
})
104+
105+
106+
async def _run(config: Config, entry_function: str, user_id: str | None = None) -> str:
107+
async with WorkflowBuilder.from_config(config=config) as builder:
108+
session_manager = await SessionManager.create(config=config,
109+
shared_builder=builder,
110+
entry_function=entry_function)
111+
try:
112+
async with session_manager.session(user_id=user_id) as session:
113+
async with session.run("abc") as runner:
114+
return await runner.result(to_type=str)
115+
finally:
116+
await session_manager.shutdown()
117+
118+
119+
@pytest.mark.parametrize("group_options", [{}, {
120+
"exclude": ["reverse"]
121+
}, {
122+
"include": ["upper"]
123+
}],
124+
ids=["no_filter", "exclude", "include"])
125+
async def test_an_available_group_function_can_be_the_entry_function(group_options):
126+
assert await _run(_config(**group_options), "text__upper") == "ABC"
127+
128+
129+
@pytest.mark.parametrize("entry_function", ["text__reverse", "text__missing", "missing"])
130+
async def test_a_function_that_is_not_available_is_still_not_found(entry_function):
131+
with pytest.raises(ValueError, match="not found"):
132+
await _run(_config(exclude=["reverse"]), entry_function)
133+
134+
135+
@pytest.mark.parametrize("group_type", ["entry_test_text_tools", "entry_test_per_user_text_tools"],
136+
ids=["shared_group", "per_user_group"])
137+
async def test_a_per_user_workflow_can_use_a_group_function_as_the_entry_function(group_type):
138+
config = _config(group_type=group_type, workflow_type="entry_test_per_user_echo_workflow")
139+
assert await _run(config, "text__upper", user_id="user-1") == "ABC"
140+
141+
142+
class NestedToolsConfig(FunctionGroupBaseConfig, name="entry_test_nested_tools"):
143+
function_name: str
144+
145+
146+
@pytest.fixture(name="register_nested_group", scope="module")
147+
def register_nested_group_fixture():
148+
149+
@register_function_group(config_type=NestedToolsConfig)
150+
async def _nested_tools(config: NestedToolsConfig, builder: Builder):
151+
group = FunctionGroup(config=config)
152+
153+
async def nested(text: str) -> str:
154+
return text
155+
156+
group.add_function(config.function_name, nested, description="a function whose name may hold the separator")
157+
yield group
158+
159+
160+
@pytest.mark.usefixtures("register_nested_group")
161+
async def test_two_groups_making_the_same_name_available_is_an_error():
162+
"""Group `a` with `b__c` and group `a__b` with `c` both offer `a__b__c`."""
163+
config = Config.model_validate({
164+
"function_groups": {
165+
"a": {
166+
"_type": "entry_test_nested_tools", "function_name": "b__c"
167+
},
168+
"a__b": {
169+
"_type": "entry_test_nested_tools", "function_name": "c"
170+
}
171+
},
172+
"workflow": {
173+
"_type": "entry_test_echo_workflow"
174+
},
175+
})
176+
with pytest.raises(ValueError, match="ambiguous"):
177+
await _run(config, "a__b__c")
Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
"""The FastMCP server publishes the functions a group makes available as tools, with or without `include`."""
16+
17+
import pytest
18+
from fastmcp import Client
19+
from fastmcp import FastMCP
20+
21+
from nat.builder.builder import Builder
22+
from nat.builder.function import FunctionGroup
23+
from nat.builder.function_info import FunctionInfo
24+
from nat.builder.workflow_builder import WorkflowBuilder
25+
from nat.cli.register_workflow import register_function
26+
from nat.cli.register_workflow import register_function_group
27+
from nat.data_models.config import Config
28+
from nat.data_models.function import FunctionBaseConfig
29+
from nat.data_models.function import FunctionGroupBaseConfig
30+
from nat.plugins.fastmcp.server.front_end_config import FastMCPFrontEndConfig
31+
from nat.plugins.fastmcp.server.front_end_plugin_worker import FastMCPFrontEndPluginWorker
32+
33+
34+
class FastMCPTextToolsConfig(FunctionGroupBaseConfig, name="fastmcp_publish_test_text_tools"):
35+
pass
36+
37+
38+
class FastMCPEchoWorkflowConfig(FunctionBaseConfig, name="fastmcp_publish_test_echo_workflow"):
39+
pass
40+
41+
42+
@pytest.fixture(name="register_components", scope="module", autouse=True)
43+
def register_components_fixture():
44+
45+
@register_function_group(config_type=FastMCPTextToolsConfig)
46+
async def _text_tools(config: FastMCPTextToolsConfig, builder: Builder):
47+
group = FunctionGroup(config=config)
48+
49+
async def upper(text: str) -> str:
50+
return text.upper()
51+
52+
async def reverse(text: str) -> str:
53+
return text[::-1]
54+
55+
group.add_function("upper", upper, description="upper-case the text")
56+
group.add_function("reverse", reverse, description="reverse the text")
57+
yield group
58+
59+
@register_function(config_type=FastMCPEchoWorkflowConfig)
60+
async def _workflow(config, builder: Builder):
61+
62+
async def _inner(text: str) -> str:
63+
return text
64+
65+
yield FunctionInfo.from_fn(_inner, description="echo")
66+
67+
68+
@pytest.mark.parametrize("group_options, tool_names, expected",
69+
[
70+
({}, [], {"text__upper", "text__reverse", "fastmcp_publish_test_echo_workflow"}),
71+
({}, ["text"], {"text__upper", "text__reverse"}),
72+
({
73+
"exclude": ["reverse"]
74+
}, ["text"], {"text__upper"}),
75+
({
76+
"include": ["upper"]
77+
}, ["text"], {"text__upper"}),
78+
],
79+
ids=["no_filter", "group_name", "exclude", "include"])
80+
async def test_fastmcp_publishes_the_group_functions(group_options, tool_names, expected):
81+
config = Config.model_validate({
82+
"function_groups": {
83+
"text": {
84+
"_type": "fastmcp_publish_test_text_tools", **group_options
85+
}
86+
},
87+
"workflow": {
88+
"_type": "fastmcp_publish_test_echo_workflow"
89+
},
90+
})
91+
config.general.front_end = FastMCPFrontEndConfig(tool_names=tool_names)
92+
mcp = FastMCP("test")
93+
94+
async with WorkflowBuilder.from_config(config=config) as builder:
95+
await FastMCPFrontEndPluginWorker(config)._default_add_routes(mcp, builder)
96+
async with Client(mcp) as client:
97+
assert {tool.name for tool in await client.list_tools()} == expected
98+
result = await client.call_tool("text__upper", {"text": "abc"})
99+
assert "ABC" in result.content[0].text

0 commit comments

Comments
 (0)