nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.3.0a20250917__py3-none-any.whl

nat/builder/workflow_builder.py

@@ -16,7 +16,9 @@
 import dataclasses
 import inspect
 import logging
+import typing
 import warnings
+from collections.abc import Sequence
 from contextlib import AbstractAsyncContextManager
 from contextlib import AsyncExitStack
 from contextlib import asynccontextmanager
@@ -31,6 +33,7 @@ from nat.builder.context import ContextState
 from nat.builder.embedder import EmbedderProviderInfo
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.builder.function import LambdaFunction
 from nat.builder.function_info import FunctionInfo
 from nat.builder.llm import LLMProviderInfo
@@ -42,6 +45,7 @@ from nat.data_models.authentication import AuthProviderBaseConfig
 from nat.data_models.component import ComponentGroup
 from nat.data_models.component_ref import AuthenticationRef
 from nat.data_models.component_ref import EmbedderRef
+from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.data_models.component_ref import LLMRef
 from nat.data_models.component_ref import MemoryRef
@@ -52,6 +56,7 @@ from nat.data_models.config import Config
 from nat.data_models.config import GeneralConfig
 from nat.data_models.embedder import EmbedderBaseConfig
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.function import FunctionGroupBaseConfig
 from nat.data_models.function_dependencies import FunctionDependencies
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.memory import MemoryBaseConfig
@@ -85,6 +90,12 @@ class ConfiguredFunction:
     instance: Function
 
 
+@dataclasses.dataclass
+class ConfiguredFunctionGroup:
+    config: FunctionGroupBaseConfig
+    instance: FunctionGroup
+
+
 @dataclasses.dataclass
 class ConfiguredLLM:
     config: LLMBaseConfig
@@ -145,6 +156,7 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
         self._telemetry_exporters: dict[str, ConfiguredTelemetryExporter] = {}
 
         self._functions: dict[str, ConfiguredFunction] = {}
+        self._function_groups: dict[str, ConfiguredFunctionGroup] = {}
         self._workflow: ConfiguredFunction | None = None
 
         self._llms: dict[str, ConfiguredLLM] = {}
@@ -161,7 +173,9 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
         # Create a mapping to track function name -> other function names it depends on
         self.function_dependencies: dict[str, FunctionDependencies] = {}
+        self.function_group_dependencies: dict[str, FunctionDependencies] = {}
         self.current_function_building: str | None = None
+        self.current_function_group_building: str | None = None
 
     async def __aenter__(self):
 
@@ -224,12 +238,32 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
         if (self._workflow is None):
             raise ValueError("Must set a workflow before building")
 
+        # Set of all functions which are "included" by function groups
+        included_functions = set()
+        # Dictionary of function configs
+        function_configs = dict()
+        # Dictionary of function group configs
+        function_group_configs = dict()
+        # Dictionary of function instances
+        function_instances = dict()
+        # Dictionary of function group instances
+        function_group_instances = dict()
+
+        for k, v in self._function_groups.items():
+            included_functions.update(v.instance.get_included_functions().keys())
+            function_group_configs[k] = v.config
+            function_group_instances[k] = v.instance
+
+        # Function configs need to be restricted to only the functions that are not in a function group
+        for k, v in self._functions.items():
+            if k not in included_functions:
+                function_configs[k] = v.config
+                function_instances[k] = v.instance
+
         # Build the config from the added objects
         config = Config(general=self.general_config,
-                        functions={
-                            k: v.config
-                            for k, v in self._functions.items()
-                        },
+                        functions=function_configs,
+                        function_groups=function_group_configs,
                         workflow=self._workflow.config,
                         llms={
                             k: v.config
@@ -263,10 +297,8 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
         workflow = Workflow.from_entry_fn(config=config,
                                           entry_fn=entry_fn_obj,
-                                          functions={
-                                              k: v.instance
-                                              for k, v in self._functions.items()
-                                          },
+                                          functions=function_instances,
+                                          function_groups=function_group_instances,
                                           llms={
                                               k: v.instance
                                               for k, v in self._llms.items()
@@ -347,11 +379,51 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
         return ConfiguredFunction(config=config, instance=build_result)
 
+    async def _build_function_group(self, name: str, config: FunctionGroupBaseConfig) -> ConfiguredFunctionGroup:
+        """Build a function group from the provided configuration.
+
+        Args:
+            name: The name of the function group
+            config: The function group configuration
+
+        Returns:
+            ConfiguredFunctionGroup: The built function group
+
+        Raises:
+            ValueError: If the function group builder returns invalid results
+        """
+        registration = self._registry.get_function_group(type(config))
+
+        inner_builder = ChildBuilder(self)
+
+        # Build the function group - use the same wrapping pattern as _build_function
+        llms = {k: v.instance for k, v in self._llms.items()}
+        function_frameworks = detect_llm_frameworks_in_build_fn(registration)
+
+        build_fn = chain_wrapped_build_fn(registration.build_fn, llms, function_frameworks)
+
+        # Set the currently building function group so the ChildBuilder can track dependencies
+        self.current_function_group_building = config.type
+        # Empty set of dependencies for the current function group
+        self.function_group_dependencies[config.type] = FunctionDependencies()
+
+        build_result = await self._get_exit_stack().enter_async_context(build_fn(config, inner_builder))
+
+        self.function_group_dependencies[name] = inner_builder.dependencies
+
+        if not isinstance(build_result, FunctionGroup):
+            raise ValueError("Expected a FunctionGroup object to be returned from the function group builder. "
+                             f"Got {type(build_result)}")
+
+        return ConfiguredFunctionGroup(config=config, instance=build_result)
+
     @override
     async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
+        if isinstance(name, FunctionRef):
+            name = str(name)
 
-        if (name in self._functions):
-            raise ValueError(f"Function `{name}` already exists in the list of functions")
+        if (name in self._functions or name in self._function_groups):
+            raise ValueError(f"Function `{name}` already exists in the list of functions or function groups")
 
         build_result = await self._build_function(name=name, config=config)
 
@@ -360,20 +432,66 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
         return build_result.instance
 
     @override
-    def get_function(self, name: str | FunctionRef) -> Function:
+    async def add_function_group(self, name: str | FunctionGroupRef, config: FunctionGroupBaseConfig) -> FunctionGroup:
+        if isinstance(name, FunctionGroupRef):
+            name = str(name)
+
+        if (name in self._function_groups or name in self._functions):
+            raise ValueError(f"Function group `{name}` already exists in the list of function groups or functions")
+
+        # Build the function group
+        build_result = await self._build_function_group(name=name, config=config)
+
+        self._function_groups[name] = build_result
+
+        # If the function group exposes functions, add them to the global function registry
+        # If the function group exposes functions, record and add them to the registry
+        for k in build_result.instance.get_included_functions():
+            if k in self._functions:
+                raise ValueError(f"Exposed function `{k}` from group `{name}` conflicts with an existing function")
+        self._functions.update({
+            k: ConfiguredFunction(config=v.config, instance=v)
+            for k, v in build_result.instance.get_included_functions().items()
+        })
+
+        return build_result.instance
 
+    @override
+    def get_function(self, name: str | FunctionRef) -> Function:
+        if isinstance(name, FunctionRef):
+            name = str(name)
         if name not in self._functions:
             raise ValueError(f"Function `{name}` not found")
 
         return self._functions[name].instance
 
+    @override
+    def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup:
+        if isinstance(name, FunctionGroupRef):
+            name = str(name)
+        if name not in self._function_groups:
+            raise ValueError(f"Function group `{name}` not found")
+
+        return self._function_groups[name].instance
+
     @override
     def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
+        if isinstance(name, FunctionRef):
+            name = str(name)
         if name not in self._functions:
             raise ValueError(f"Function `{name}` not found")
 
         return self._functions[name].config
 
+    @override
+    def get_function_group_config(self, name: str | FunctionGroupRef) -> FunctionGroupBaseConfig:
+        if isinstance(name, FunctionGroupRef):
+            name = str(name)
+        if name not in self._function_groups:
+            raise ValueError(f"Function group `{name}` not found")
+
+        return self._function_groups[name].config
+
     @override
     async def set_workflow(self, config: FunctionBaseConfig) -> Function:
 
@@ -403,16 +521,59 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
     @override
     def get_function_dependencies(self, fn_name: str | FunctionRef) -> FunctionDependencies:
+        if isinstance(fn_name, FunctionRef):
+            fn_name = str(fn_name)
         return self.function_dependencies[fn_name]
 
     @override
-    def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str):
+    def get_function_group_dependencies(self, fn_name: str | FunctionGroupRef) -> FunctionDependencies:
+        if isinstance(fn_name, FunctionGroupRef):
+            fn_name = str(fn_name)
+        return self.function_group_dependencies[fn_name]
+
+    @override
+    def get_tools(self,
+                  tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
+                  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
+        tools = []
+        seen = set()
+        for n in tool_names:
+            is_function_group_ref = isinstance(n, FunctionGroupRef)
+            if isinstance(n, FunctionRef) or is_function_group_ref:
+                n = str(n)
+            if n in seen:
+                raise ValueError(f"Function or Function Group `{n}` already seen")
+            seen.add(n)
+            if n not in self._function_groups:
+                # the passed tool name is probably a function
+                if is_function_group_ref:
+                    raise ValueError(f"Function group `{n}` not found in the list of function groups")
+                tools.append(self.get_tool(n, wrapper_type))
+                continue
+
+            # Using the registry, get the tool wrapper for the requested framework
+            tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
+
+            current_function_group = self._function_groups[n]
+
+            # walk through all functions in the function group -- guaranteed to not be fallible
+            for fn_name, fn_instance in current_function_group.instance.get_accessible_functions().items():
+                try:
+                    # Wrap in the correct wrapper and add to tools list
+                    tools.append(tool_wrapper_reg.build_fn(fn_name, fn_instance, self))
+                except Exception:
+                    logger.error("Error fetching tool `%s`", fn_name, exc_info=True)
+                    raise
+
+        return tools
 
+    @override
+    def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str):
+        if isinstance(fn_name, FunctionRef):
+            fn_name = str(fn_name)
         if fn_name not in self._functions:
             raise ValueError(f"Function `{fn_name}` not found in list of functions")
-
         fn = self._functions[fn_name]
-
         try:
             # Using the registry, get the tool wrapper for the requested framework
             tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
@@ -892,12 +1053,15 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
                 # Instantiate a memory client
                 elif component_instance.component_group == ComponentGroup.MEMORY:
                     await self.add_memory_client(component_instance.name, component_instance.config)
-            # Instantiate a object store client
+                # Instantiate a object store client
                 elif component_instance.component_group == ComponentGroup.OBJECT_STORES:
                     await self.add_object_store(component_instance.name, component_instance.config)
                 # Instantiate a retriever client
                 elif component_instance.component_group == ComponentGroup.RETRIEVERS:
                     await self.add_retriever(component_instance.name, component_instance.config)
+                # Instantiate a function group
+                elif component_instance.component_group == ComponentGroup.FUNCTION_GROUPS:
+                    await self.add_function_group(component_instance.name, component_instance.config)
                 # Instantiate a function
                 elif component_instance.component_group == ComponentGroup.FUNCTIONS:
                     # If the function is the root, set it as the workflow later
@@ -956,6 +1120,10 @@ class ChildBuilder(Builder):
     async def add_function(self, name: str, config: FunctionBaseConfig) -> Function:
         return await self._workflow_builder.add_function(name, config)
 
+    @override
+    async def add_function_group(self, name: str, config: FunctionGroupBaseConfig) -> FunctionGroup:
+        return await self._workflow_builder.add_function_group(name, config)
+
     @override
     def get_function(self, name: str) -> Function:
         # If a function tries to get another function, we assume it uses it
@@ -965,10 +1133,23 @@ class ChildBuilder(Builder):
 
         return fn
 
+    @override
+    def get_function_group(self, name: str) -> FunctionGroup:
+        # If a function tries to get a function group, we assume it uses it
+        function_group = self._workflow_builder.get_function_group(name)
+
+        self._dependencies.add_function_group(name)
+
+        return function_group
+
     @override
     def get_function_config(self, name: str) -> FunctionBaseConfig:
         return self._workflow_builder.get_function_config(name)
 
+    @override
+    def get_function_group_config(self, name: str) -> FunctionGroupBaseConfig:
+        return self._workflow_builder.get_function_group_config(name)
+
     @override
     async def set_workflow(self, config: FunctionBaseConfig) -> Function:
         return await self._workflow_builder.set_workflow(config)
@@ -982,7 +1163,19 @@ class ChildBuilder(Builder):
         return self._workflow_builder.get_workflow_config()
 
     @override
-    def get_tool(self, fn_name: str, wrapper_type: LLMFrameworkEnum | str):
+    def get_tools(self,
+                  tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
+                  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
+        tools = self._workflow_builder.get_tools(tool_names, wrapper_type)
+        for tool_name in tool_names:
+            if tool_name in self._workflow_builder._function_groups:
+                self._dependencies.add_function_group(tool_name)
+            else:
+                self._dependencies.add_function(tool_name)
+        return tools
+
+    @override
+    def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str):
         # If a function tries to get another function as a tool, we assume it uses it
         fn = self._workflow_builder.get_tool(fn_name, wrapper_type)
 
@@ -1111,3 +1304,7 @@ class ChildBuilder(Builder):
     @override
     def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
         return self._workflow_builder.get_function_dependencies(fn_name)
+
+    @override
+    def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
+        return self._workflow_builder.get_function_group_dependencies(fn_name)

nat/cli/commands/optimize.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+import asyncio
+import logging
+from pathlib import Path
+
+import click
+
+from nat.data_models.optimizer import OptimizerRunConfig
+from nat.profiler.parameter_optimization.optimizer_runtime import optimize_config
+
+logger = logging.getLogger(__name__)
+
+
+@click.group(name=__name__, invoke_without_command=True, help="Optimize a workflow with the specified dataset.")
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=True,
+    help="A JSON/YAML file that sets the parameters for the workflow and evaluation.",
+)
+@click.option(
+    "--dataset",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=False,
+    help="A json file with questions and ground truth answers. This will override the dataset path in the config file.",
+)
+@click.option(
+    "--result_json_path",
+    type=str,
+    default="$",
+    help=("A JSON path to extract the result from the workflow. Use this when the workflow returns "
+          "multiple objects or a dictionary. For example, '$.output' will extract the 'output' field "
+          "from the result."),
+)
+@click.option(
+    "--endpoint",
+    type=str,
+    default=None,
+    help="Use endpoint for running the workflow. Example: http://localhost:8000/generate",
+)
+@click.option(
+    "--endpoint_timeout",
+    type=int,
+    default=300,
+    help="HTTP response timeout in seconds. Only relevant if endpoint is specified.",
+)
+@click.pass_context
+def optimizer_command(ctx, **kwargs) -> None:
+    """ Optimize workflow with the specified dataset"""
+    pass
+
+
+async def run_optimizer(config: OptimizerRunConfig):
+    await optimize_config(config)
+
+
+@optimizer_command.result_callback(replace=True)
+def run_optimizer_callback(
+    processors,  # pylint: disable=unused-argument
+    *,
+    config_file: Path,
+    dataset: Path,
+    result_json_path: str,
+    endpoint: str,
+    endpoint_timeout: int,
+):
+    """Run the optimizer with the provided config file and dataset."""
+    config = OptimizerRunConfig(
+        config_file=config_file,
+        dataset=dataset,
+        result_json_path=result_json_path,
+        endpoint=endpoint,
+        endpoint_timeout=endpoint_timeout,
+    )
+
+    asyncio.run(run_optimizer(config))

nat/cli/commands/workflow/templates/config.yml.j2

@@ -1,5 +1,4 @@
 general:
-  use_uvloop: true
   logging:
     console:
       _type: console

nat/cli/commands/workflow/workflow_commands.py

@@ -171,6 +171,9 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
         workflow_dir (str): The directory to create the workflow package.
         description (str): Description to pre-popluate the workflow docstring.
     """
+    # Fail fast with Click's standard exit code (2) for bad params.
+    if not workflow_name or not workflow_name.strip():
+        raise click.BadParameter("Workflow name cannot be empty.")  # noqa: TRY003
     try:
         # Get the repository root
         try:
@@ -217,15 +220,13 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
         else:
             install_cmd = ['pip', 'install', '-e', str(new_workflow_dir)]
 
-        config_source = configs_dir / 'config.yml'
-
         # 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'{workflow_name}_function.py',
             '__init__.py.j2': base_dir / '__init__.py',
-            'config.yml.j2': config_source,
+            'config.yml.j2': configs_dir / 'config.yml',
         }
 
         # Render templates
@@ -246,10 +247,6 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
             with open(output_path, 'w', encoding="utf-8") as f:
                 f.write(content)
 
-        # Create symlink for config.yml
-        config_link = new_workflow_dir / 'configs' / 'config.yml'
-        os.symlink(config_source, config_link)
-
         # Create symlinks for config and data directories
         config_dir_source = configs_dir
         config_dir_link = new_workflow_dir / 'configs'

nat/cli/entrypoint.py

@@ -34,6 +34,7 @@ from .commands.configure.configure import configure_command
 from .commands.evaluate import eval_command
 from .commands.info.info import info_command
 from .commands.object_store.object_store import object_store_command
+from .commands.optimize import optimizer_command
 from .commands.registry.registry import registry_command
 from .commands.sizing.sizing import sizing
 from .commands.start import start_command
@@ -108,6 +109,7 @@ cli.add_command(uninstall_command, name="uninstall")
 cli.add_command(validate_command, name="validate")
 cli.add_command(workflow_command, name="workflow")
 cli.add_command(sizing, name="sizing")
+cli.add_command(optimizer_command, name="optimize")
 cli.add_command(object_store_command, name="object-store")
 
 # Aliases

nat/cli/register_workflow.py

@@ -27,6 +27,8 @@ from nat.cli.type_registry import EvaluatorRegisteredCallableT
 from nat.cli.type_registry import FrontEndBuildCallableT
 from nat.cli.type_registry import FrontEndRegisteredCallableT
 from nat.cli.type_registry import FunctionBuildCallableT
+from nat.cli.type_registry import FunctionGroupBuildCallableT
+from nat.cli.type_registry import FunctionGroupRegisteredCallableT
 from nat.cli.type_registry import FunctionRegisteredCallableT
 from nat.cli.type_registry import LLMClientBuildCallableT
 from nat.cli.type_registry import LLMClientRegisteredCallableT
@@ -60,6 +62,7 @@ from nat.data_models.embedder import EmbedderBaseConfigT
 from nat.data_models.evaluator import EvaluatorBaseConfigT
 from nat.data_models.front_end import FrontEndConfigT
 from nat.data_models.function import FunctionConfigT
+from nat.data_models.function import FunctionGroupConfigT
 from nat.data_models.llm import LLMBaseConfigT
 from nat.data_models.memory import MemoryBaseConfigT
 from nat.data_models.object_store import ObjectStoreBaseConfigT
@@ -155,10 +158,7 @@ def register_function(config_type: type[FunctionConfigT],
 
         context_manager_fn = asynccontextmanager(fn)
 
-        if framework_wrappers is None:
-            framework_wrappers_list: list[str] = []
-        else:
-            framework_wrappers_list = list(framework_wrappers)
+        framework_wrappers_list = list(framework_wrappers or [])
 
         discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
                                                                 component_type=ComponentEnum.FUNCTION)
@@ -177,6 +177,40 @@ def register_function(config_type: type[FunctionConfigT],
     return register_function_inner
 
 
+def register_function_group(config_type: type[FunctionGroupConfigT],
+                            framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
+    """
+    Register a function group with optional framework_wrappers for automatic profiler hooking.
+    Function groups share configuration/resources across multiple functions.
+    """
+
+    def register_function_group_inner(
+        fn: FunctionGroupBuildCallableT[FunctionGroupConfigT]
+    ) -> FunctionGroupRegisteredCallableT[FunctionGroupConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredFunctionGroupInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        framework_wrappers_list = list(framework_wrappers or [])
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.FUNCTION_GROUP)
+
+        GlobalTypeRegistry.get().register_function_group(
+            RegisteredFunctionGroupInfo(
+                full_type=config_type.full_type,
+                config_type=config_type,
+                build_fn=context_manager_fn,
+                framework_wrappers=framework_wrappers_list,
+                discovery_metadata=discovery_metadata,
+            ))
+
+        return context_manager_fn
+
+    return register_function_group_inner
+
+
 def register_llm_provider(config_type: type[LLMBaseConfigT]):
 
     def register_llm_provider_inner(