nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/cli/commands/start.py

@@ -0,0 +1,257 @@
+# 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 functools
+import logging
+import typing
+from collections.abc import Callable
+from pathlib import Path
+
+import click
+from pydantic_core import SchemaValidator
+
+from nat.cli.cli_utils.config_override import load_and_override_config
+from nat.cli.type_registry import GlobalTypeRegistry
+from nat.cli.type_registry import RegisteredFrontEndInfo
+from nat.data_models.config import Config
+from nat.utils.data_models.schema_validator import validate_schema
+from nat.utils.type_utils import DecomposedType
+
+logger = logging.getLogger(__name__)
+
+
+class StartCommandGroup(click.Group):
+
+    def __init__(
+        self,
+        name: str | None = None,
+        invoke_without_command: bool = False,
+        no_args_is_help: bool | None = None,
+        subcommand_metavar: str | None = None,
+        chain: bool = False,
+        result_callback: Callable[..., typing.Any] | None = None,
+        **attrs: typing.Any,
+    ):
+        super().__init__(name=name,
+                         invoke_without_command=invoke_without_command,
+                         no_args_is_help=no_args_is_help,
+                         subcommand_metavar=subcommand_metavar,
+                         chain=chain,
+                         result_callback=result_callback,
+                         **attrs)
+
+        self._commands: dict[str, click.Command] | None = None
+        self._registered_front_ends: dict[str, RegisteredFrontEndInfo] = {}
+
+    def _build_params(self, front_end: RegisteredFrontEndInfo) -> list[click.Parameter]:
+
+        params: list[click.Parameter] = []
+
+        # First two are always the config file and override
+        params.append(
+            click.Option(param_decls=["--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.")))
+        params.append(
+            click.Option(
+                param_decls=["--override"],
+                type=(str, str),
+                multiple=True,
+                help="Override config values using dot notation (e.g., --override llms.nim_llm.temperature 0.7)"))
+
+        fields = front_end.config_type.model_fields
+        for name, field in fields.items():
+
+            if (name in ("override", "config_file")):
+                raise ValueError(
+                    "Cannot have a field named 'override' or 'config_file' in the front end config. These are reserved."
+                )
+
+            # Skip init-only fields since we dont want to set them in the constructor. Must check for False explicitly
+            if (field.init == False):  # noqa: E712, pylint: disable=singleton-comparison
+                continue
+
+            if (field.annotation is None):
+                raise ValueError(f"Field {name} has no type annotation. Types are required for Front End Plugins.")
+
+            # Decompose the type into its origin and arguments
+            decomposed_type = DecomposedType(field.annotation)
+
+            param_decls = [f"--{name}"]
+            multiple = False
+
+            # Remove any optional types
+            while (decomposed_type.is_optional):
+                decomposed_type = decomposed_type.get_optional_type()
+
+            if (decomposed_type.is_union):
+                raise ValueError(f"Invalid field '{name}'.Unions are only supported for optional parameters.")
+
+            # Handle the types
+            # Literal[...] -> map to click.Choice([...])
+            if (decomposed_type.origin is typing.Literal):
+                # typing.get_args returns the literal values; ensure they are strings for Click
+                literal_values = [str(v) for v in decomposed_type.args]
+                param_type = click.Choice(literal_values)
+
+            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)):
+                if (len(decomposed_type.args) == 1):
+                    inner = DecomposedType(decomposed_type.args[0])
+                    # Support containers of Literal values -> multiple Choice
+                    if (inner.origin is typing.Literal):
+                        literal_values = [str(v) for v in inner.args]
+                        param_type = click.Choice(literal_values)
+                    else:
+                        param_type = inner.root
+                else:
+                    param_type = None
+
+                multiple = True
+            else:
+                param_type = decomposed_type.root
+
+            if (field.alias is not None):
+                param_decls = [f"--{field.alias}", f"{name}"]
+
+            params.append(
+                click.Option(param_decls=param_decls,
+                             type=param_type,
+                             required=False,
+                             multiple=multiple,
+                             help=field.description))
+
+        return params
+
+    def _load_commands(self) -> dict[str, click.Command]:
+
+        if (self._commands is not None):
+            return self._commands
+
+        from nat.runtime.loader import PluginTypes
+        from nat.runtime.loader import discover_and_register_plugins
+
+        # Only load front ends here for performance. Ensures a responsive CLI
+        discover_and_register_plugins(PluginTypes.FRONT_END)
+
+        all_front_ends = GlobalTypeRegistry.get().get_registered_front_ends()
+
+        self._commands = {}
+
+        for front_end in all_front_ends:
+
+            registered_front_end = GlobalTypeRegistry.get().get_front_end(config_type=front_end.config_type)
+
+            # Build the command parameters
+            params: list[click.Parameter] = self._build_params(registered_front_end)
+            help_msg = f"Run a NAT workflow using the {registered_front_end.local_name} front end."
+
+            cmd = click.Command(name=registered_front_end.local_name,
+                                params=params,
+                                help=help_msg,
+                                callback=functools.partial(click.pass_context(self.invoke_subcommand),
+                                                           cmd_name=front_end.local_name))
+
+            self._registered_front_ends[front_end.local_name] = registered_front_end
+            self._commands[front_end.local_name] = cmd
+
+        return self._commands
+
+    def invoke_subcommand(self,
+                          ctx: click.Context,
+                          cmd_name: str,
+                          config_file: Path,
+                          override: tuple[tuple[str, str], ...],
+                          **kwargs) -> int | None:
+
+        from nat.runtime.loader import PluginTypes
+        from nat.runtime.loader import discover_and_register_plugins
+
+        if (config_file is None):
+            raise click.ClickException("No config file provided.")
+
+        # Here we need to ensure all objects are loaded before we try to create the config object
+        discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
+
+        logger.info("Starting NAT from config file: '%s'", config_file)
+
+        config_dict = load_and_override_config(config_file, override)
+
+        # Get the front end for the command
+        front_end: RegisteredFrontEndInfo = self._registered_front_ends[cmd_name]
+
+        config = validate_schema(config_dict, Config)
+
+        # Override default front end config with values from the config file for serverless execution modes.
+        # Check that we have the right kind of front end
+        if (not isinstance(config.general.front_end, front_end.config_type)):
+
+            # Set the front end config
+            config.general.front_end = front_end.config_type()
+
+        front_end_config = config.general.front_end
+
+        # Iterate over the parameters and set them in the config
+        for param, value in kwargs.items():
+
+            # Skip default values so we dont overwrite the config
+            if (ctx.get_parameter_source(param) == click.core.ParameterSource.DEFAULT):
+                continue
+
+            setattr(front_end_config, param, value)
+
+        # Validate the config once more to ensure that all parameters are set correctly
+        schema_validator = SchemaValidator(schema=front_end_config.__pydantic_core_schema__)
+        schema_validator.validate_python(front_end_config.__dict__)
+
+        try:
+
+            async def run_plugin():
+
+                # From the config, get the registered front end plugin
+                front_end_info = GlobalTypeRegistry.get().get_front_end(config_type=type(front_end_config))
+
+                # Create the front end plugin
+                async with front_end_info.build_fn(front_end_config, config) as front_end_plugin:
+
+                    # Run the front end plugin
+                    await front_end_plugin.run()
+
+            return asyncio.run(run_plugin())
+
+        except Exception as e:
+            logger.error("Failed to initialize workflow")
+            raise click.ClickException(str(e)) from e
+
+    def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None:
+
+        return self._load_commands().get(cmd_name)
+
+    def list_commands(self, ctx: click.Context) -> list[str]:
+        return sorted(self._load_commands().keys())
+
+
+@click.command(name=__name__,
+               invoke_without_command=False,
+               help="Run a NAT workflow using a front end configuration.",
+               cls=StartCommandGroup)
+@click.pass_context
+def start_command(ctx: click.Context, **kwargs) -> None:
+    """Run a NAT workflow using a front end configuration."""
+    pass

nat/cli/commands/uninstall.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 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 contextlib import AsyncExitStack
+
+import click
+
+logger = logging.getLogger(__name__)
+
+
+async def uninstall_packages(packages: list[dict[str, str]]) -> None:
+
+    from nat.cli.type_registry import GlobalTypeRegistry
+    from nat.registry_handlers.schemas.package import PackageNameVersionList
+    from nat.runtime.loader import PluginTypes
+    from nat.runtime.loader import discover_and_register_plugins
+    from nat.settings.global_settings import GlobalSettings
+
+    discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
+
+    registry = GlobalTypeRegistry.get()
+
+    config_dict = {"channels": {"uninstall_local": {"_type": "local"}}}
+    registry_config = GlobalSettings.get().model_validate(config_dict)
+    local_registry_config = registry_config.channels.get("uninstall_local", None)
+
+    if (local_registry_config is None):
+        logger.error("Channel runtime instance not found.")
+
+    try:
+        package_name_list = PackageNameVersionList(**{"packages": packages})
+    except Exception as e:
+        logger.exception("Error validating package format: %s", e)
+        return
+
+    async with AsyncExitStack() as stack:
+        registry_handler_info = registry.get_registry_handler(type(local_registry_config))
+        registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(local_registry_config))
+        await stack.enter_async_context(registry_handler.remove(packages=package_name_list))
+
+
+@click.group(name=__name__, invoke_without_command=True, help=("Uninstall plugin packages from the local environment."))
+@click.argument("packages", type=str)
+def uninstall_command(packages: str) -> None:
+    """
+    Uninstall plugin packages from the local environment.
+    """
+
+    packages = packages.split()
+    packages_versions = []
+    for package in packages:
+        package_dict = {}
+        package_version = package.split("==")
+        if (len(package_version) == 1):
+            package_dict["name"] = package_version[0]
+            msg = f"No package version provided for '{package_version[0]}'."
+            logger.warning(msg)
+        elif (len(package_version) == 2):
+            package_dict["name"] = package_version[0]
+            package_dict["version"] = package_version[1]
+        else:
+            msg = f"Invalid input: '{package}'"
+            logger.error(msg)
+        if (package_dict):
+            packages_versions.append(package_dict)
+
+    asyncio.run(uninstall_packages(packages=packages_versions))

nat/cli/commands/validate.py

@@ -0,0 +1,47 @@
+# 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.
+
+from pathlib import Path
+
+import click
+
+
+@click.command()
+@click.option("--config_file",
+              type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+              required=True,
+              help="Configuration file to validate")
+def validate_command(config_file: Path):
+    """Validate a configuration file"""
+    # load function level dependencies
+    from io import StringIO
+
+    from nat.runtime.loader import load_config
+
+    try:
+        click.echo(f"Validating configuration file: {config_file}")
+        config = load_config(config_file)
+        click.echo(click.style("✓ Configuration file is valid!", fg="green"))
+
+        stream = StringIO()
+
+        config.print_summary(stream=stream)
+
+        click.echo_via_pager(stream.getvalue())
+    except Exception as e:
+        click.echo(click.style("✗ Validation failed!\n\nError:", fg="red"))
+
+        click.echo(click.style(e, fg="red"))
+        raise click.ClickException(str(e)) from e

nat/cli/commands/workflow/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 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.

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

@@ -0,0 +1,17 @@
+functions:
+  current_datetime:
+    _type: current_datetime
+  {{python_safe_workflow_name}}:
+    _type: {{python_safe_workflow_name}}
+    prefix: "Hello:"
+
+llms:
+  nim_llm:
+    _type: nim
+    model_name: meta/llama-3.1-70b-instruct
+    temperature: 0.0
+
+workflow:
+  _type: react_agent
+  llm_name: nim_llm
+  tool_names: [current_datetime, {{python_safe_workflow_name}}]

nat/cli/commands/workflow/templates/pyproject.toml.j2

@@ -0,0 +1,25 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+{% if editable %}requires = ["setuptools >= 64", "setuptools-scm>=8"]
+
+[tool.setuptools_scm]
+# NAT uses the --first-parent flag to avoid tags from previous releases which have been merged into the develop branch
+# from causing an unexpected version change. This can be safely removed if developing outside of the NAT repository.
+git_describe_command = "git describe --long --first-parent"
+root = "{{ rel_path_to_repo_root}}"{% else %}requires = ["setuptools >= 64"]{% endif %}
+
+[project]
+name = "{{ package_name }}"
+{% if editable %}dynamic = ["version"]{% else %}version = "0.1.0"{% endif %}
+dependencies = [
+  "{{ nat_dependency }}",
+]
+requires-python = ">=3.11,<3.14"
+description = "Custom NeMo Agent Toolkit Workflow"
+classifiers = ["Programming Language :: Python"]
+
+{% if editable %}[tool.uv.sources]
+nvidia-nat = { path = "{{ rel_path_to_repo_root}}", editable = true }{% endif %}
+
+[project.entry-points.'nat.components']
+{{ package_name }} = "{{ package_name }}.register"

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

@@ -0,0 +1,4 @@
+# flake8: noqa
+
+# 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

@@ -0,0 +1,50 @@
+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
+
+logger = logging.getLogger(__name__)
+
+
+class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"):
+    """
+    {{ workflow_description }}
+    """
+    prefix: str = Field(default="Echo:", description="Prefix to add before the echoed text.")
+
+
+@register_function(config_type={{ workflow_class_name }}, 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.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 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 logging
+
+import click
+
+from nat.cli.commands.workflow.workflow_commands import create_command
+from nat.cli.commands.workflow.workflow_commands import delete_command
+from nat.cli.commands.workflow.workflow_commands import reinstall_command
+
+logger = logging.getLogger(__name__)
+
+
+@click.group(name=__name__, invoke_without_command=False, help="Interact with templated workflows.")
+def workflow_command(**kwargs):
+    """
+    Interact with templated workflows.
+    """
+    pass
+
+
+workflow_command.add_command(create_command, name="create")
+workflow_command.add_command(delete_command, "delete")
+workflow_command.add_command(reinstall_command, "reinstall")