nvidia-nat 1.2.1rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

nat/cli/commands/workflow/workflow_commands.py

@@ -27,6 +27,77 @@ 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.
+
+    Args:
+        versioned: Whether to include the version in the dependency string
+
+    Returns:
+        str: The dependency string to use in pyproject.toml
+    """
+    # Assume the default dependency is LangChain/LangGraph
+    dependency = "nvidia-nat[langchain]"
+
+    if not versioned:
+        logger.debug("Using unversioned NAT dependency: %s", dependency)
+        return dependency
+
+    version = _get_nat_version()
+    if version is None:
+        logger.debug("Could not detect NAT version, using unversioned dependency: %s", dependency)
+        return dependency
+
+    dependency += f"~={version}"
+    logger.debug("Using NAT dependency: %s", dependency)
+    return dependency
+
+
 class PackageError(Exception):
     pass
 
@@ -66,7 +137,7 @@ def find_package_root(package_name: str) -> Path | None:
         try:
             info = json.loads(direct_url)
         except json.JSONDecodeError:
-            logger.error("Malformed direct_url.json for package: %s", package_name)
+            logger.exception("Malformed direct_url.json for package: %s", package_name)
             return None
 
         if not info.get("dir_info", {}).get("editable"):
@@ -130,7 +201,6 @@ def get_workflow_path_from_name(workflow_name: str):
     default="NAT function template. Please update the description.",
     help="""A description of the component being created. Will be used to populate the docstring and will describe the
          component when inspecting installed components using 'nat info component'""")
-# pylint: disable=missing-param-doc
 def create_command(workflow_name: str, install: bool, workflow_dir: str, description: str):
     """
     Create a new NAT workflow using templates.
@@ -141,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:
@@ -166,12 +239,17 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
             click.echo(f"Workflow '{workflow_name}' already exists.")
             return
 
+        base_dir = new_workflow_dir / 'src' / package_name
+
+        configs_dir = base_dir / 'configs'
+        data_dir = base_dir / 'data'
+
         # Create directory structure
-        (new_workflow_dir / 'src' / package_name).mkdir(parents=True)
+        base_dir.mkdir(parents=True)
         # Create config directory
-        (new_workflow_dir / 'src' / package_name / 'configs').mkdir(parents=True)
-        # Create package level configs directory
-        (new_workflow_dir / 'configs').mkdir(parents=True)
+        configs_dir.mkdir(parents=True)
+        # Create data directory
+        data_dir.mkdir(parents=True)
 
         # Initialize Jinja2 environment
         env = Environment(loader=FileSystemLoader(str(template_dir)))
@@ -181,25 +259,30 @@ 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")
+
+        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': new_workflow_dir / 'src' / package_name / 'register.py',
-            'workflow.py.j2': new_workflow_dir / 'src' / package_name / f'{workflow_name}_function.py',
-            '__init__.py.j2': new_workflow_dir / 'src' / package_name / '__init__.py',
-            'config.yml.j2': new_workflow_dir / 'src' / package_name / 'configs' / 'config.yml',
+            'register.py.j2': base_dir / 'register.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',
         }
 
         # 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",
-            'workflow_description': description
+            'workflow_description': description,
+            'nat_dependency': _get_nat_dependency()
         }
 
         for template_name, output_path in files_to_render.items():
@@ -208,10 +291,13 @@ 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_source = new_workflow_dir / 'src' / package_name / 'configs' / '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'
+        data_dir_source = data_dir
+        data_dir_link = new_workflow_dir / 'data'
+        os.symlink(config_dir_source, config_dir_link)
+        os.symlink(data_dir_source, data_dir_link)
 
         if install:
             # Install the new package without changing directories
@@ -226,7 +312,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
 
         click.echo(f"Workflow '{workflow_name}' created successfully in '{new_workflow_dir}'.")
     except Exception as e:
-        logger.exception("An error occurred while creating the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while creating the workflow: %s", e)
         click.echo(f"An error occurred while creating the workflow: {e}")
 
 
@@ -262,7 +348,7 @@ def reinstall_command(workflow_name):
 
         click.echo(f"Workflow '{workflow_name}' reinstalled successfully.")
     except Exception as e:
-        logger.exception("An error occurred while reinstalling the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while reinstalling the workflow: %s", e)
         click.echo(f"An error occurred while reinstalling the workflow: {e}")
 
 
@@ -309,7 +395,7 @@ def delete_command(workflow_name: str):
 
         click.echo(f"Workflow '{workflow_name}' deleted successfully.")
     except Exception as e:
-        logger.exception("An error occurred while deleting the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while deleting the workflow: %s", e)
         click.echo(f"An error occurred while deleting the workflow: {e}")
 
 

nat/cli/entrypoint.py

@@ -29,10 +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
@@ -40,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
 
 
@@ -107,11 +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")
 
 
 @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,
@@ -298,7 +313,7 @@ class RegisteredPackage(BaseModel):
     discovery_metadata: DiscoveryMetadata
 
 
-class TypeRegistry:  # pylint: disable=too-many-public-methods
+class TypeRegistry:
 
     def __init__(self) -> None:
         # Telemetry Exporters
@@ -313,6 +328,9 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         # 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:  # pylint: disable=too-many-public-methods
 
         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):
@@ -588,8 +650,8 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         except KeyError as err:
             raise KeyError(
                 f"An invalid Embedder config and wrapper combination was supplied. Config: `{config_type}`, "
-                "Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Embedder client but "
-                "there is no registered conversion from that Embedder provider to LLM framework: {wrapper_type}. "
+                f"Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Embedder client but "
+                f"there is no registered conversion from that Embedder provider to LLM framework: {wrapper_type}. "
                 "Please provide an Embedder configuration from one of the following providers: "
                 f"{set(self._embedder_client_provider_to_framework.keys())}") from err
 
@@ -703,8 +765,8 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         except KeyError as err:
             raise KeyError(
                 f"An invalid Retriever config and wrapper combination was supplied. Config: `{config_type}`, "
-                "Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Retriever client but "
-                "there is no registered conversion from that Retriever provider to LLM framework: {wrapper_type}. "
+                f"Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Retriever client but "
+                f"there is no registered conversion from that Retriever provider to LLM framework: {wrapper_type}. "
                 "Please provide a Retriever configuration from one of the following providers: "
                 f"{set(self._retriever_client_provider_to_framework.keys())}") from err
 
@@ -779,7 +841,7 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         self._registration_changed()
 
-    def get_infos_by_type(self, component_type: ComponentEnum) -> dict:  # pylint: disable=R0911
+    def get_infos_by_type(self, component_type: ComponentEnum) -> dict:
 
         if component_type == ComponentEnum.FRONT_END:
             return self._registered_front_end_infos
@@ -790,6 +852,9 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         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
 
@@ -849,12 +914,14 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         raise ValueError(f"Supplied an unsupported component type {component_type}")
 
-    def get_registered_types_by_component_type(  # pylint: disable=R0911
-            self, component_type: ComponentEnum) -> list[str]:
+    def get_registered_types_by_component_type(self, component_type: ComponentEnum) -> list[str]:
 
         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)
 
@@ -925,8 +992,7 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
             if (short_names[key.local_name] == 1):
                 type_list.append((key.local_name, key.config_type))
 
-        # pylint: disable=consider-alternative-union-syntax
-        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]):
 
@@ -945,6 +1011,9 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         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