nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

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/start.py

@@ -111,7 +111,7 @@ class StartCommandGroup(click.Group):
             elif (issubclass(decomposed_type.root, Path)):
                 param_type = click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path)
 
-            elif (issubclass(decomposed_type.root, (list, tuple, set))):
+            elif (issubclass(decomposed_type.root, list | tuple | set)):
                 if (len(decomposed_type.args) == 1):
                     inner = DecomposedType(decomposed_type.args[0])
                     # Support containers of Literal values -> multiple Choice

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

@@ -1,16 +1,17 @@
-general:
-  use_uvloop: true
-  logging:
-    console:
-      _type: console
-      level: WARN
+functions:
+  current_datetime:
+    _type: current_datetime
+  {{python_safe_workflow_name}}:
+    _type: {{python_safe_workflow_name}}
+    prefix: "Hello:"
 
-  front_end:
-    _type: fastapi
-
-  front_end:
-    _type: console
+llms:
+  nim_llm:
+    _type: nim
+    model_name: meta/llama-3.1-70b-instruct
+    temperature: 0.0
 
 workflow:
-  _type: {{workflow_name}}
-  parameter: default_value
+  _type: react_agent
+  llm_name: nim_llm
+  tool_names: [current_datetime, {{python_safe_workflow_name}}]

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 .{{package_name}} import {{ python_safe_workflow_name }}_function

nat/cli/commands/workflow/templates/workflow.py.j2

@@ -3,6 +3,7 @@ import logging
 from pydantic import Field
 
 from nat.builder.builder import Builder
+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
@@ -12,25 +13,38 @@ logger = logging.getLogger(__name__)
 
 class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"):
     """
-    {{workflow_description}}
+    {{ workflow_description }}
     """
-    # Add your custom configuration parameters here
-    parameter: str = Field(default="default_value", description="Notional description for this parameter")
-
-
-@register_function(config_type={{ workflow_class_name }})
-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.")
+    prefix: str = Field(default="Echo:", description="Prefix to add before the echoed text.")
+
+
+@register_function(config_type={{ workflow_class_name }}, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def {{ python_safe_workflow_name }}_function(config: {{ workflow_class_name }}, builder: Builder):
+    """
+    Registers a function (addressable via `{{ workflow_name }}` in the configuration).
+    This registration ensures a static mapping of the function type, `{{ workflow_name }}`, to the `{{ workflow_class_name }}` configuration object.
+
+    Args:
+        config ({{ workflow_class_name }}): The configuration for the function.
+        builder (Builder): The builder object.
+
+    Returns:
+        FunctionInfo: The function info object for the function.
+    """
+
+    # Define the function that will be registered.
+    async def _echo(text: str) -> str:
+        """
+        Takes a text input and echoes back with a pre-defined prefix.
+
+        Args:
+            text (str): The text to echo back.
+
+        Returns:
+            str: The text with the prefix.
+        """
+        return f"{config.prefix} {text}"
+
+    # The callable is wrapped in a FunctionInfo object.
+    # The description parameter is used to describe the function.
+    yield FunctionInfo.from_fn(_echo, description=_echo.__doc__)

nat/cli/commands/workflow/workflow_commands.py

@@ -27,6 +27,50 @@ from jinja2 import FileSystemLoader
 logger = logging.getLogger(__name__)
 
 
+def _get_nat_version() -> str | None:
+    """
+    Get the current NAT version.
+
+    Returns:
+        str: The NAT version intended for use in a dependency string.
+        None: If the NAT version is not found.
+    """
+    from nat.cli.entrypoint import get_version
+
+    current_version = get_version()
+    if current_version == "unknown":
+        return None
+
+    version_parts = current_version.split(".")
+    if len(version_parts) < 3:
+        # If the version somehow doesn't have three parts, return the full version
+        return current_version
+
+    patch = version_parts[2]
+    try:
+        # If the patch is a number, keep only the major and minor parts
+        # Useful for stable releases and adheres to semantic versioning
+        _ = int(patch)
+        digits_to_keep = 2
+    except ValueError:
+        # If the patch is not a number, keep all three digits
+        # Useful for pre-release versions (and nightly builds)
+        digits_to_keep = 3
+
+    return ".".join(version_parts[:digits_to_keep])
+
+
+def _is_nat_version_prerelease() -> bool:
+    """
+    Check if the NAT version is a prerelease.
+    """
+    version = _get_nat_version()
+    if version is None:
+        return False
+
+    return len(version.split(".")) >= 3
+
+
 def _get_nat_dependency(versioned: bool = True) -> str:
     """
     Get the NAT dependency string with version.
@@ -44,16 +88,12 @@ def _get_nat_dependency(versioned: bool = True) -> str:
         logger.debug("Using unversioned NAT dependency: %s", dependency)
         return dependency
 
-    # Get the current NAT version
-    from nat.cli.entrypoint import get_version
-    current_version = get_version()
-    if current_version == "unknown":
-        logger.warning("Could not detect NAT version, using unversioned dependency")
+    version = _get_nat_version()
+    if version is None:
+        logger.debug("Could not detect NAT version, using unversioned dependency: %s", dependency)
         return dependency
 
-    # Extract major.minor (e.g., "1.2.3" -> "1.2")
-    major_minor = ".".join(current_version.split(".")[:2])
-    dependency += f"~={major_minor}"
+    dependency += f"~={version}"
     logger.debug("Using NAT dependency: %s", dependency)
     return dependency
 
@@ -171,6 +211,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:
@@ -216,23 +259,25 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
             install_cmd = ['uv', 'pip', 'install', '-e', str(new_workflow_dir)]
         else:
             install_cmd = ['pip', 'install', '-e', str(new_workflow_dir)]
+            if _is_nat_version_prerelease():
+                install_cmd.insert(2, "--pre")
 
-        config_source = configs_dir / 'config.yml'
+        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'{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': config_source,
+            'config.yml.j2': configs_dir / 'config.yml',
         }
 
         # Render templates
         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",
@@ -246,10 +291,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'
@@ -313,7 +354,8 @@ def reinstall_command(workflow_name):
 
 @click.command()
 @click.argument('workflow_name')
-def delete_command(workflow_name: str):
+@click.option('-y', '--yes', "yes_flag", is_flag=True, default=False, help='Do not prompt for confirmation.')
+def delete_command(workflow_name: str, yes_flag: bool):
     """
     Delete a NAT workflow and uninstall its package.
 
@@ -321,7 +363,7 @@ def delete_command(workflow_name: str):
         workflow_name (str): The name of the workflow to delete.
     """
     try:
-        if not click.confirm(f"Are you sure you want to delete the workflow '{workflow_name}'?"):
+        if not yes_flag and not click.confirm(f"Are you sure you want to delete the workflow '{workflow_name}'?"):
             click.echo("Workflow deletion cancelled.")
             return
         editable = get_repo_root() is not None

nat/cli/entrypoint.py

@@ -29,11 +29,16 @@ import time
 
 import click
 import nest_asyncio
+from dotenv import load_dotenv
+
+from nat.utils.log_levels import LOG_LEVELS
 
 from .commands.configure.configure import configure_command
 from .commands.evaluate import eval_command
 from .commands.info.info import info_command
+from .commands.mcp.mcp import mcp_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
@@ -41,23 +46,21 @@ from .commands.uninstall import uninstall_command
 from .commands.validate import validate_command
 from .commands.workflow.workflow import workflow_command
 
+# Load environment variables from .env file, if it exists
+load_dotenv()
+
 # Apply at the beginning of the file to avoid issues with asyncio
 nest_asyncio.apply()
 
-# Define log level choices
-LOG_LEVELS = {
-    'DEBUG': logging.DEBUG,
-    'INFO': logging.INFO,
-    'WARNING': logging.WARNING,
-    'ERROR': logging.ERROR,
-    'CRITICAL': logging.CRITICAL
-}
-
 
 def setup_logging(log_level: str):
     """Configure logging with the specified level"""
     numeric_level = LOG_LEVELS.get(log_level.upper(), logging.INFO)
-    logging.basicConfig(level=numeric_level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    logging.basicConfig(
+        level=numeric_level,
+        format="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
     return numeric_level
 
 
@@ -108,12 +111,13 @@ 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")
+cli.add_command(mcp_command, name="mcp")
 
 # Aliases
 cli.add_command(start_command.get_command(None, "console"), name="run")  # type: ignore
 cli.add_command(start_command.get_command(None, "fastapi"), name="serve")  # type: ignore
-cli.add_command(start_command.get_command(None, "mcp"), name="mcp")  # type: ignore
 
 
 @cli.result_callback()

nat/cli/main.py

@@ -30,6 +30,9 @@ def run_cli():
     import os
     import sys
 
+    # Suppress warnings from transformers
+    os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
     parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
     if (parent_dir not in sys.path):

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(

nat/cli/type_registry.py

@@ -37,6 +37,7 @@ from nat.builder.embedder import EmbedderProviderInfo
 from nat.builder.evaluator import EvaluatorInfo
 from nat.builder.front_end import FrontEndBase
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.builder.function_base import FunctionBase
 from nat.builder.function_info import FunctionInfo
 from nat.builder.llm import LLMProviderInfo
@@ -55,6 +56,8 @@ from nat.data_models.front_end import FrontEndBaseConfig
 from nat.data_models.front_end import FrontEndConfigT
 from nat.data_models.function import FunctionBaseConfig
 from nat.data_models.function import FunctionConfigT
+from nat.data_models.function import FunctionGroupBaseConfig
+from nat.data_models.function import FunctionGroupConfigT
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.llm import LLMBaseConfigT
 from nat.data_models.logging import LoggingBaseConfig
@@ -85,6 +88,7 @@ EmbedderProviderBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncI
 EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIterator[EvaluatorInfo]]
 FrontEndBuildCallableT = Callable[[FrontEndConfigT, Config], AsyncIterator[FrontEndBase]]
 FunctionBuildCallableT = Callable[[FunctionConfigT, Builder], AsyncIterator[FunctionInfo | Callable | FunctionBase]]
+FunctionGroupBuildCallableT = Callable[[FunctionGroupConfigT, Builder], AsyncIterator[FunctionGroup]]
 TTCStrategyBuildCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AsyncIterator[StrategyBase]]
 LLMClientBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[typing.Any]]
 LLMProviderBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[LLMProviderInfo]]
@@ -106,6 +110,7 @@ EvaluatorRegisteredCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], Abs
 FrontEndRegisteredCallableT = Callable[[FrontEndConfigT, Config], AbstractAsyncContextManager[FrontEndBase]]
 FunctionRegisteredCallableT = Callable[[FunctionConfigT, Builder],
                                        AbstractAsyncContextManager[FunctionInfo | Callable | FunctionBase]]
+FunctionGroupRegisteredCallableT = Callable[[FunctionGroupConfigT, Builder], AbstractAsyncContextManager[FunctionGroup]]
 TTCStrategyRegisterCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AbstractAsyncContextManager[StrategyBase]]
 LLMClientRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
 LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
@@ -178,6 +183,16 @@ class RegisteredFunctionInfo(RegisteredInfo[FunctionBaseConfig]):
     framework_wrappers: list[str] = Field(default_factory=list)
 
 
+class RegisteredFunctionGroupInfo(RegisteredInfo[FunctionGroupBaseConfig]):
+    """
+    Represents a registered function group. Function groups are collections of functions that share configuration
+    and resources.
+    """
+
+    build_fn: FunctionGroupRegisteredCallableT = Field(repr=False)
+    framework_wrappers: list[str] = Field(default_factory=list)
+
+
 class RegisteredLLMProviderInfo(RegisteredInfo[LLMBaseConfig]):
     """
     Represents a registered LLM provider. LLM Providers are the operators of the LLMs. i.e. NIMs, OpenAI, Anthropic,
@@ -313,6 +328,9 @@ class TypeRegistry:
         # Functions
         self._registered_functions: dict[type[FunctionBaseConfig], RegisteredFunctionInfo] = {}
 
+        # Function Groups
+        self._registered_function_groups: dict[type[FunctionGroupBaseConfig], RegisteredFunctionGroupInfo] = {}
+
         # LLMs
         self._registered_llm_provider_infos: dict[type[LLMBaseConfig], RegisteredLLMProviderInfo] = {}
         self._llm_client_provider_to_framework: dict[type[LLMBaseConfig], dict[str, RegisteredLLMClientInfo]] = {}
@@ -478,6 +496,50 @@ class TypeRegistry:
 
         return list(self._registered_functions.values())
 
+    def register_function_group(self, registration: RegisteredFunctionGroupInfo):
+        """Register a function group with the type registry.
+
+        Args:
+            registration: The function group registration information
+
+        Raises:
+            ValueError: If a function group with the same config type is already registered
+        """
+        if (registration.config_type in self._registered_function_groups):
+            raise ValueError(
+                f"A function group with the same config type `{registration.config_type}` has already been "
+                "registered.")
+
+        self._registered_function_groups[registration.config_type] = registration
+
+        self._registration_changed()
+
+    def get_function_group(self, config_type: type[FunctionGroupBaseConfig]) -> RegisteredFunctionGroupInfo:
+        """Get a registered function group by its config type.
+
+        Args:
+            config_type: The function group configuration type
+
+        Returns:
+            RegisteredFunctionGroupInfo: The registered function group information
+
+        Raises:
+            KeyError: If no function group is registered for the given config type
+        """
+        try:
+            return self._registered_function_groups[config_type]
+        except KeyError as err:
+            raise KeyError(f"Could not find a registered function group for config `{config_type}`. "
+                           f"Registered configs: {set(self._registered_function_groups.keys())}") from err
+
+    def get_registered_function_groups(self) -> list[RegisteredInfo[FunctionGroupBaseConfig]]:
+        """Get all registered function groups.
+
+        Returns:
+            list[RegisteredInfo[FunctionGroupBaseConfig]]: List of all registered function groups
+        """
+        return list(self._registered_function_groups.values())
+
     def register_llm_provider(self, info: RegisteredLLMProviderInfo):
 
         if (info.config_type in self._registered_llm_provider_infos):
@@ -790,6 +852,9 @@ class TypeRegistry:
         if component_type == ComponentEnum.FUNCTION:
             return self._registered_functions
 
+        if component_type == ComponentEnum.FUNCTION_GROUP:
+            return self._registered_function_groups
+
         if component_type == ComponentEnum.TOOL_WRAPPER:
             return self._registered_tool_wrappers
 
@@ -854,6 +919,9 @@ class TypeRegistry:
         if component_type == ComponentEnum.FUNCTION:
             return [i.static_type() for i in self._registered_functions]
 
+        if component_type == ComponentEnum.FUNCTION_GROUP:
+            return [i.static_type() for i in self._registered_function_groups]
+
         if component_type == ComponentEnum.TOOL_WRAPPER:
             return list(self._registered_tool_wrappers)
 
@@ -924,7 +992,7 @@ class TypeRegistry:
             if (short_names[key.local_name] == 1):
                 type_list.append((key.local_name, key.config_type))
 
-        return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
+        return typing.Union[*tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
 
     def compute_annotation(self, cls: type[TypedBaseModelT]):
 
@@ -943,6 +1011,9 @@ class TypeRegistry:
         if issubclass(cls, FunctionBaseConfig):
             return self._do_compute_annotation(cls, self.get_registered_functions())
 
+        if issubclass(cls, FunctionGroupBaseConfig):
+            return self._do_compute_annotation(cls, self.get_registered_function_groups())
+
         if issubclass(cls, LLMBaseConfig):
             return self._do_compute_annotation(cls, self.get_registered_llm_providers())
 

nat/control_flow/register.py

@@ -0,0 +1,20 @@
+# 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.
+
+# flake8: noqa
+
+# Import any control flows which need to be automatically registered here
+from . import sequential_executor
+from .router_agent import register