nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/cli/commands/start.py

@@ -0,0 +1,246 @@
+# 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 aiq.cli.cli_utils.config_override import load_and_override_config
+from aiq.cli.type_registry import GlobalTypeRegistry
+from aiq.cli.type_registry import RegisteredFrontEndInfo
+from aiq.data_models.config import AIQConfig
+from aiq.utils.data_models.schema_validator import validate_schema
+from aiq.utils.type_utils import DecomposedType
+
+logger = logging.getLogger(__name__)
+
+
+class StartCommandGroup(click.Group):
+
+    # pylint: disable=too-many-positional-arguments
+    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
+            if (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):
+                    param_type = decomposed_type.args[0]
+                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 aiq.runtime.loader import PluginTypes
+        from aiq.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 an AIQ Toolkit 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 aiq.runtime.loader import PluginTypes
+        from aiq.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 AIQ Toolkit 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, AIQConfig)
+
+        # 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", exc_info=True)
+            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 an AIQ Toolkit workflow using a front end configuration.",
+               cls=StartCommandGroup)
+@click.pass_context
+def start_command(ctx: click.Context, **kwargs) -> None:
+    """Run an AIQ Toolkit workflow using a front end configuration."""
+    pass

aiq/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 aiq.cli.type_registry import GlobalTypeRegistry
+    from aiq.registry_handlers.schemas.package import PackageNameVersionList
+    from aiq.runtime.loader import PluginTypes
+    from aiq.runtime.loader import discover_and_register_plugins
+    from aiq.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, exc_info=True)
+        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))

aiq/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 aiq.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

aiq/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.

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

@@ -0,0 +1,16 @@
+general:
+  use_uvloop: true
+  logging:
+    console:
+      _type: console
+      level: WARN
+
+  front_end:
+    _type: fastapi
+
+  front_end:
+    _type: console
+
+workflow:
+  _type: {{workflow_name}}
+  parameter: default_value

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

@@ -0,0 +1,22 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+{% if editable %}requires = ["setuptools >= 64", "setuptools-scm>=8"]
+
+[tool.setuptools_scm]
+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 = [
+  "aiqtoolkit[langchain]",
+]
+requires-python = ">=3.11,<3.13"
+description = "Custom AIQ Toolkit Workflow"
+classifiers = ["Programming Language :: Python"]
+
+{% if editable %}[tool.uv.sources]
+aiqtoolkit = { path = "{{ rel_path_to_repo_root}}", editable = true }{% endif %}
+
+[project.entry-points.'aiq.components']
+{{ package_name }} = "{{ package_name }}.register"

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

@@ -0,0 +1,5 @@
+# pylint: disable=unused-import
+# flake8: noqa
+
+# Import any tools which need to be automatically registered here
+from {{package_name}} import {{workflow_name}}_function

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

@@ -0,0 +1,36 @@
+import logging
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"):
+    """
+    {{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.")

aiq/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 aiq.cli.commands.workflow.workflow_commands import create_command
+from aiq.cli.commands.workflow.workflow_commands import delete_command
+from aiq.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")