nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc2__py3-none-any.whl

nat/cli/commands/object_store/__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/object_store/object_store.py

@@ -0,0 +1,227 @@
+# 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 importlib
+import logging
+import mimetypes
+import time
+from pathlib import Path
+
+import click
+
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.object_store import ObjectStoreBaseConfig
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+logger = logging.getLogger(__name__)
+
+STORE_CONFIGS = {
+    "s3": {
+        "module": "nat.plugins.s3.object_store", "config_class": "S3ObjectStoreClientConfig"
+    },
+    "mysql": {
+        "module": "nat.plugins.mysql.object_store", "config_class": "MySQLObjectStoreClientConfig"
+    },
+    "redis": {
+        "module": "nat.plugins.redis.object_store", "config_class": "RedisObjectStoreClientConfig"
+    }
+}
+
+
+def get_object_store_config(**kwargs) -> ObjectStoreBaseConfig:
+    """Process common object store arguments and return the config class"""
+    store_type = kwargs.pop("store_type")
+    config = STORE_CONFIGS[store_type]
+    module = importlib.import_module(config["module"])
+    config_class = getattr(module, config["config_class"])
+    return config_class(**kwargs)
+
+
+async def upload_file(object_store: ObjectStore, file_path: Path, key: str):
+    """
+    Upload a single file to object store.
+
+    Args:
+        object_store: The object store instance to use.
+        file_path: The path to the file to upload.
+        key: The key to upload the file to.
+    """
+    try:
+        data = await asyncio.to_thread(file_path.read_bytes)
+
+        item = ObjectStoreItem(data=data,
+                               content_type=mimetypes.guess_type(str(file_path))[0],
+                               metadata={
+                                   "original_filename": file_path.name,
+                                   "file_size": str(len(data)),
+                                   "file_extension": file_path.suffix,
+                                   "upload_timestamp": str(int(time.time()))
+                               })
+
+        # Upload using upsert to allow overwriting
+        await object_store.upsert_object(key, item)
+        click.echo(f"✅ Uploaded: {file_path.name} -> {key}")
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to upload {file_path.name}:\n{e}") from e
+
+
+def object_store_command_decorator(async_func):
+    """
+    Decorator that handles the common object store command pattern.
+
+    The decorated function should take (store: ObjectStore, kwargs) as parameters
+    and return an exit code (0 for success).
+    """
+
+    @click.pass_context
+    def wrapper(ctx: click.Context, **kwargs):
+        config = ctx.obj["store_config"]
+
+        async def work():
+            async with WorkflowBuilder() as builder:
+                await builder.add_object_store(name="store", config=config)
+                store = await builder.get_object_store_client("store")
+                return await async_func(store, **kwargs)
+
+        try:
+            exit_code = asyncio.run(work())
+        except Exception as e:
+            raise click.ClickException(f"Command failed: {e}") from e
+        if exit_code != 0:
+            raise click.ClickException(f"Command failed with exit code {exit_code}")
+        return exit_code
+
+    return wrapper
+
+
+@click.command(name="upload", help="Upload a directory to an object store.")
+@click.argument("local_dir",
+                type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
+                required=True)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def upload_command(store: ObjectStore, local_dir: Path, **_kwargs):
+    """
+    Upload a directory to an object store.
+
+    Args:
+        local_dir: The local directory to upload.
+        store: The object store to use.
+        _kwargs: Additional keyword arguments.
+    """
+    try:
+        click.echo(f"📁 Processing directory: {local_dir}")
+        file_count = 0
+
+        # Process each file recursively
+        for file_path in local_dir.rglob("*"):
+            if file_path.is_file():
+                key = file_path.relative_to(local_dir).as_posix()
+                await upload_file(store, file_path, key)
+                file_count += 1
+
+        click.echo(f"✅ Directory uploaded successfully! {file_count} files uploaded.")
+        return 0
+
+    except Exception as e:
+        raise click.ClickException(f"❌ Failed to upload directory {local_dir}:\n  {e}") from e
+
+
+@click.command(name="delete", help="Delete files from an object store.")
+@click.argument("keys", type=str, required=True, nargs=-1)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def delete_command(store: ObjectStore, keys: list[str], **_kwargs):
+    """
+    Delete files from an object store.
+
+    Args:
+        store: The object store to use.
+        keys: The keys to delete.
+        _kwargs: Additional keyword arguments.
+    """
+    deleted_count = 0
+    failed_count = 0
+    for key in keys:
+        try:
+            await store.delete_object(key)
+            click.echo(f"✅ Deleted: {key}")
+            deleted_count += 1
+        except Exception as e:
+            click.echo(f"❌ Failed to delete {key}: {e}")
+            failed_count += 1
+
+    click.echo(f"✅ Deletion completed! {deleted_count} keys deleted. {failed_count} keys failed to delete.")
+    return 0 if failed_count == 0 else 1
+
+
+@click.group(name="object-store", invoke_without_command=False, help="Manage object store operations.")
+def object_store_command(**_kwargs):
+    """Manage object store operations including uploading files and directories."""
+    pass
+
+
+def register_object_store_commands():
+
+    @click.group(name="s3", invoke_without_command=False, help="S3 object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--endpoint-url", type=str, help="S3 endpoint URL")
+    @click.option("--access-key", type=str, help="S3 access key")
+    @click.option("--secret-key", type=str, help="S3 secret key")
+    @click.option("--region", type=str, help="S3 region")
+    @click.pass_context
+    def s3(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="s3", **kwargs)
+
+    @click.group(name="mysql", invoke_without_command=False, help="MySQL object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="MySQL host")
+    @click.option("--port", type=int, help="MySQL port")
+    @click.option("--db", type=str, help="MySQL database name")
+    @click.option("--username", type=str, help="MySQL username")
+    @click.option("--password", type=str, help="MySQL password")
+    @click.pass_context
+    def mysql(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="mysql", **kwargs)
+
+    @click.group(name="redis", invoke_without_command=False, help="Redis object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="Redis host")
+    @click.option("--port", type=int, help="Redis port")
+    @click.option("--db", type=int, help="Redis db")
+    @click.pass_context
+    def redis(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="redis", **kwargs)
+
+    commands = {"s3": s3, "mysql": mysql, "redis": redis}
+
+    for store_type, config in STORE_CONFIGS.items():
+        try:
+            importlib.import_module(config["module"])
+            command = commands[store_type]
+            object_store_command.add_command(command, name=store_type)
+            command.add_command(upload_command, name="upload")
+            command.add_command(delete_command, name="delete")
+        except ImportError:
+            pass
+
+
+register_object_store_commands()

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/registry/publish.py

@@ -40,7 +40,7 @@ async def publish_artifact(registry_handler_config: RegistryHandlerBaseConfig, p
         try:
             artifact = build_artifact(package_root=package_root)
         except Exception as e:
-            logger.exception("Error building artifact: %s", e, exc_info=True)
+            logger.exception("Error building artifact: %s", e)
             return
         await stack.enter_async_context(registry_handler.publish(artifact=artifact))
 
@@ -82,7 +82,7 @@ def publish(channel: str, config_file: str, package_root: str) -> None:
             logger.error("Publish channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(publish_artifact(registry_handler_config=publish_channel_config, package_root=package_root))

nat/cli/commands/registry/pull.py

@@ -66,7 +66,7 @@ async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, pack
             validated_packages = PullRequestPackages(packages=package_list)
 
         except Exception as e:
-            logger.exception("Error processing package names: %s", e, exc_info=True)
+            logger.exception("Error processing package names: %s", e)
             return
 
         await stack.enter_async_context(registry_handler.pull(packages=validated_packages))
@@ -112,7 +112,7 @@ def pull(channel: str, config_file: str, packages: str) -> None:
             logger.error("Pull channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(pull_artifact(pull_channel_config, packages))

nat/cli/commands/registry/remove.py

@@ -41,7 +41,7 @@ async def remove_artifact(registry_handler_config: RegistryHandlerBaseConfig, pa
         try:
             package_name_list = PackageNameVersionList(**{"packages": packages})
         except Exception as e:
-            logger.exception("Invalid package format: '%s'", e, exc_info=True)
+            logger.exception("Invalid package format: '%s'", e)
 
         await stack.enter_async_context(registry_handler.remove(packages=package_name_list))
 
@@ -102,7 +102,7 @@ def remove(channel: str, config_file: str, packages: str) -> None:
             logger.error("Remove channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(remove_artifact(registry_handler_config=remove_channel_config, packages=packages_versions))

nat/cli/commands/registry/search.py

@@ -29,14 +29,13 @@ from nat.utils.data_models.schema_validator import validate_yaml
 logger = logging.getLogger(__name__)
 
 
-async def search_artifacts(  # pylint: disable=R0917
-        registry_handler_config: RegistryHandlerBaseConfig,
-        query: str,
-        search_fields: list[SearchFields],
-        visualize: bool,
-        component_types: list[ComponentEnum],
-        save_path: str | None = None,
-        n_results: int = 10) -> None:
+async def search_artifacts(registry_handler_config: RegistryHandlerBaseConfig,
+                           query: str,
+                           search_fields: list[SearchFields],
+                           visualize: bool,
+                           component_types: list[ComponentEnum],
+                           save_path: str | None = None,
+                           n_results: int = 10) -> None:
 
     from nat.cli.type_registry import GlobalTypeRegistry
     from nat.registry_handlers.schemas.search import SearchQuery
@@ -116,14 +115,13 @@ async def search_artifacts(  # pylint: disable=R0917
     required=False,
     help=("The component types to include in search."),
 )
-def search(  # pylint: disable=R0917
-        config_file: str,
-        channel: str,
-        fields: list[str],
-        query: str,
-        component_types: list[ComponentEnum],
-        n_results: int,
-        output_path: str) -> None:
+def search(config_file: str,
+           channel: str,
+           fields: list[str],
+           query: str,
+           component_types: list[ComponentEnum],
+           n_results: int,
+           output_path: str) -> None:
     """
     Search for NAT artifacts with the specified configuration.
     """
@@ -142,7 +140,7 @@ def search(  # pylint: disable=R0917
             logger.error("Search channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(

nat/cli/commands/start.py

@@ -35,7 +35,6 @@ logger = logging.getLogger(__name__)
 
 class StartCommandGroup(click.Group):
 
-    # pylint: disable=too-many-positional-arguments
     def __init__(
         self,
         name: str | None = None,
@@ -103,12 +102,24 @@ class StartCommandGroup(click.Group):
                 raise ValueError(f"Invalid field '{name}'.Unions are only supported for optional parameters.")
 
             # Handle the types
-            if (issubclass(decomposed_type.root, Path)):
+            # 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))):
+            elif (issubclass(decomposed_type.root, list | tuple | set)):
                 if (len(decomposed_type.args) == 1):
-                    param_type = decomposed_type.args[0]
+                    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
 
@@ -225,7 +236,7 @@ class StartCommandGroup(click.Group):
             return asyncio.run(run_plugin())
 
         except Exception as e:
-            logger.error("Failed to initialize workflow", exc_info=True)
+            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:

nat/cli/commands/uninstall.py

@@ -44,7 +44,7 @@ async def uninstall_packages(packages: list[dict[str, str]]) -> None:
     try:
         package_name_list = PackageNameVersionList(**{"packages": packages})
     except Exception as e:
-        logger.exception("Error validating package format: %s", e, exc_info=True)
+        logger.exception("Error validating package format: %s", e)
         return
 
     async with AsyncExitStack() as stack:

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/pyproject.toml.j2

@@ -3,6 +3,9 @@ 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]
@@ -11,7 +14,7 @@ name = "{{ package_name }}"
 dependencies = [
   "{{ nat_dependency }}",
 ]
-requires-python = ">=3.11,<3.13"
+requires-python = ">=3.11,<3.14"
 description = "Custom NeMo Agent Toolkit Workflow"
 classifiers = ["Programming Language :: Python"]
 

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

@@ -1,5 +1,4 @@
-# pylint: disable=unused-import
 # 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__)