aiqtoolkit 1.1.0a20250515__py3-none-any.whl → 1.1.0a20251020__py3-none-any.whl

aiq/cli/commands/registry/search.py

@@ -1,155 +0,0 @@
-# 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
-from pathlib import Path
-
-import click
-
-from aiq.data_models.component import AIQComponentEnum
-from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
-from aiq.registry_handlers.schemas.search import SearchFields
-from aiq.registry_handlers.schemas.status import StatusEnum
-from aiq.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[AIQComponentEnum],
-        save_path: str | None = None,
-        n_results: int = 10) -> None:
-
-    from aiq.cli.type_registry import GlobalTypeRegistry
-    from aiq.registry_handlers.schemas.search import SearchQuery
-
-    registry = GlobalTypeRegistry.get()
-
-    async with AsyncExitStack() as stack:
-
-        registry_handler_info = registry.get_registry_handler(type(registry_handler_config))
-        registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(registry_handler_config))
-
-        if (len(component_types) == 0):
-            component_types = [t.value for t in AIQComponentEnum]
-
-        query = SearchQuery(query=query, fields=search_fields, top_k=n_results, component_types=component_types)
-
-        search_response = await stack.enter_async_context(registry_handler.search(query=query))
-
-        if (search_response.status.status == StatusEnum.SUCCESS):
-            if (visualize):
-                registry_handler.visualize_search_results(search_response=search_response)
-            if (save_path is not None):
-                registry_handler.save_search_results(search_response=search_response, save_path=save_path)
-
-
-@click.group(name=__name__, invoke_without_command=True, help="Search for AIQ Toolkit artifacts from remote registry.")
-@click.option(
-    "--config_file",
-    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
-    callback=validate_yaml,
-    required=False,
-    help=("A JSON/YAML file that sets the parameters for the workflow."),
-)
-@click.option(
-    "-c",
-    "--channel",
-    type=str,
-    required=True,
-    help=("The remote registry channel to use when pulling the AIQ Toolkit artifact."),
-)
-@click.option(
-    "-o",
-    "--output_path",
-    type=str,
-    required=False,
-    help=("Path to save search results."),
-)
-@click.option(
-    "-f",
-    "--fields",
-    multiple=True,
-    type=click.Choice([e.value for e in SearchFields], case_sensitive=False),
-    required=False,
-    help=("The fields to include in the search."),
-)
-@click.option(
-    "-q",
-    "--query",
-    type=str,
-    required=True,
-    help=("The query string."),
-)
-@click.option(
-    "-n",
-    "--n_results",
-    type=int,
-    required=False,
-    default=10,
-    help=("Number of search results to return."),
-)
-@click.option(
-    "-t",
-    "--types",
-    "component_types",
-    multiple=True,
-    type=click.Choice([e.value for e in AIQComponentEnum], case_sensitive=False),
-    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[AIQComponentEnum],
-        n_results: int,
-        output_path: str) -> None:
-    """
-    Search for AIQ Toolkit artifacts with the specified configuration.
-    """
-
-    from aiq.settings.global_settings import GlobalSettings
-
-    settings = GlobalSettings().get()
-
-    if (config_file is not None):
-        settings = settings.override_settings(config_file)
-
-    try:
-        search_channel_config = settings.channels.get(channel)
-
-        if (search_channel_config is None):
-            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)
-        return
-
-    asyncio.run(
-        search_artifacts(registry_handler_config=search_channel_config,
-                         query=query,
-                         component_types=component_types,
-                         search_fields=fields,
-                         visualize=True,
-                         save_path=output_path,
-                         n_results=n_results))

aiq/cli/commands/start.py

@@ -1,250 +0,0 @@
-# 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.MultiCommand):
-
-    # 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)
-
-        # Check that we have the right kind of front end
-        if (not isinstance(config.general.front_end, front_end.config_type)):
-            logger.warning(
-                "The front end type in the config file (%s) does not match the command name (%s). "
-                "Overwriting the config file front end.",
-                config.general.front_end.type,
-                cmd_name)
-
-            # 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

@@ -1,83 +0,0 @@
-# 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 an AIQ Toolkit plugin packages from the local environment."))
-@click.argument("packages", type=str)
-def uninstall_command(packages: str) -> None:
-    """
-    Uninstall AIQ Toolkit 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

@@ -1,47 +0,0 @@
-# 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

@@ -1,14 +0,0 @@
-# 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

@@ -1,16 +0,0 @@
-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

@@ -1,22 +0,0 @@
-[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

@@ -1,5 +0,0 @@
-# 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

@@ -1,36 +0,0 @@
-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:
-        print("Function exited early!")
-    finally:
-        print("Cleaning up {{ workflow_name }} workflow.")

aiq/cli/commands/workflow/workflow.py

@@ -1,37 +0,0 @@
-# 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")