aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/cli/commands/info/list_components.py

@@ -0,0 +1,129 @@
+# 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
+
+from aiq.data_models.component import AIQComponentEnum
+from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
+from aiq.registry_handlers.schemas.search import SearchFields
+
+logger = logging.getLogger(__name__)
+
+
+async def search_artifacts(  # pylint: disable=R0917
+        registry_handler_config: RegistryHandlerBaseConfig,
+        component_types: list[AIQComponentEnum],
+        visualize: bool,
+        query: str,
+        num_results: int,
+        query_fields: list[SearchFields],
+        save_path: str | None) -> 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]
+
+        if (len(query_fields) == 0):
+            query_fields = (SearchFields.ALL, )
+
+        query = SearchQuery(query=query, fields=query_fields, top_k=num_results, component_types=component_types)
+        search_response = await stack.enter_async_context(registry_handler.search(query=query))
+
+        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="List the locally registered AgentIQ components.")
+@click.option(
+    "-t",
+    "--types",
+    "component_types",
+    multiple=True,
+    type=click.Choice([e.value for e in AIQComponentEnum], case_sensitive=False),
+    required=False,
+    help=("Filter the search by AgentIQ component type."),
+)
+@click.option(
+    "-o",
+    "--output_path",
+    type=str,
+    required=False,
+    help=("Path to save search results."),
+)
+@click.option(
+    "-q",
+    "--query",
+    type=str,
+    default="",
+    required=False,
+    help=("The query string."),
+)
+@click.option(
+    "-n",
+    "--num_results",
+    type=int,
+    default=-1,
+    required=False,
+    help=("Number of results to return."),
+)
+@click.option(
+    "-f",
+    "--fields",
+    multiple=True,
+    type=click.Choice([e.value for e in SearchFields], case_sensitive=False),
+    required=False,
+    help=("Fields used when applying query."),
+)
+def list_components(fields: list[SearchFields],
+                    query: str,
+                    num_results: int,
+                    component_types: list[AIQComponentEnum],
+                    output_path: str | None = None) -> None:
+
+    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.ALL)
+
+    config_dict = {"channels": {"list_components": {"_type": "local"}}}
+    registry_config = GlobalSettings.get().model_validate(config_dict)
+    local_registry_config = registry_config.channels.get("list_components", None)
+    if (local_registry_config is None):
+        logger.error("Channel runtime instance not found.")
+
+    asyncio.run(
+        search_artifacts(local_registry_config,
+                         query=query,
+                         num_results=num_results,
+                         query_fields=fields,
+                         component_types=component_types,
+                         visualize=True,
+                         save_path=output_path))

aiq/cli/commands/registry/__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/registry/publish.py

@@ -0,0 +1,88 @@
+# 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.registry_handler import RegistryHandlerBaseConfig
+from aiq.utils.data_models.schema_validator import validate_yaml
+
+logger = logging.getLogger(__name__)
+
+
+async def publish_artifact(registry_handler_config: RegistryHandlerBaseConfig, package_root: str) -> None:
+
+    from aiq.cli.type_registry import GlobalTypeRegistry
+    from aiq.registry_handlers.package_utils import build_aiq_artifact
+
+    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))
+        try:
+            artifact = build_aiq_artifact(package_root=package_root)
+        except Exception as e:
+            logger.exception("Error building artifact: %s", e, exc_info=True)
+            return
+        await stack.enter_async_context(registry_handler.publish(artifact=artifact))
+
+
+@click.group(name=__name__,
+             invoke_without_command=True,
+             help=("Publish local AgentIQ artifacts to a remote "
+                   "registry from package repository."))
+@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 YAML file to override configured channel settings."),
+)
+@click.option(
+    "-c",
+    "--channel",
+    type=str,
+    required=True,
+    help=("The remote registry channel to use when publishing the AgentIQ artifact."),
+)
+@click.argument("package_root", type=str)
+def publish(channel: str, config_file: str, package_root: str) -> None:
+    """
+    Publish AgentIQ 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:
+        publish_channel_config = settings.channels.get(channel)
+
+        if (publish_channel_config is 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)
+        return
+
+    asyncio.run(publish_artifact(registry_handler_config=publish_channel_config, package_root=package_root))

aiq/cli/commands/registry/pull.py

@@ -0,0 +1,118 @@
+# 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.registry_handler import RegistryHandlerBaseConfig
+from aiq.utils.data_models.schema_validator import validate_yaml
+
+logger = logging.getLogger(__name__)
+
+
+async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, packages: list[str]) -> None:
+
+    from aiq.cli.type_registry import GlobalTypeRegistry
+    from aiq.registry_handlers.schemas.package import PackageNameVersion
+    from aiq.registry_handlers.schemas.pull import PullPackageWhl
+    from aiq.registry_handlers.schemas.pull import PullRequestPackages
+
+    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))
+
+        try:
+            package_list = []
+            for package in packages:
+
+                package_data = {}
+
+                assert len(package) > 0, f"Supplied invalid package '{package}'."
+
+                if package[:-4] == ".whl":
+                    package_data["whl_path"] = package
+                    package_list.append(PullPackageWhl(**package_data))
+                else:
+                    package_split = package.split("==")
+
+                    assert len(package_split) in (1, 2), f"Supplied invalid package '{package}'."
+
+                    package_data["name"] = package_split[0]
+
+                    if (package_split == 2):
+                        package_data["version"] = package_split[1]
+
+                    package_list.append(PackageNameVersion(**package_data))
+
+            validated_packages = PullRequestPackages(packages=package_list)
+
+        except Exception as e:
+            logger.exception("Error processing package names: %s", e, exc_info=True)
+            return
+
+        await stack.enter_async_context(registry_handler.pull(packages=validated_packages))
+
+
+@click.group(name=__name__,
+             invoke_without_command=True,
+             help=("Pull AgentIQ artifacts from a remote registry "
+                   "by package name."))
+@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 YAML file to override the channel settings."),
+)
+@click.option(
+    "-c",
+    "--channel",
+    type=str,
+    required=True,
+    help=("The remote registry channel to use when pulling the AgentIQ artifact."),
+)
+@click.argument("packages", type=str)
+def pull(channel: str, config_file: str, packages: str) -> None:
+    """
+    Pull AgentIQ artifacts from a remote registry channel.
+    """
+
+    from aiq.settings.global_settings import GlobalSettings
+
+    packages = packages.split()
+
+    settings = GlobalSettings().get()
+
+    if (config_file is not None):
+        settings = settings.override_settings(config_file)
+
+    try:
+        pull_channel_config = settings.channels.get(channel)
+
+        if (pull_channel_config is 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)
+        return
+
+    asyncio.run(pull_artifact(pull_channel_config, packages))

aiq/cli/commands/registry/registry.py

@@ -0,0 +1,36 @@
+# 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.registry.publish import publish
+from aiq.cli.commands.registry.pull import pull
+from aiq.cli.commands.registry.remove import remove
+from aiq.cli.commands.registry.search import search
+
+logger = logging.getLogger(__name__)
+
+
+@click.group(name=__name__, invoke_without_command=False, help="Utility to configure AgentIQ remote registry channels.")
+def registry_command(**kwargs):
+    pass
+
+
+registry_command.add_command(publish, "publish")
+registry_command.add_command(pull, "pull")
+registry_command.add_command(remove, "remove")
+registry_command.add_command(search, "search")

aiq/cli/commands/registry/remove.py

@@ -0,0 +1,108 @@
+# 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.registry_handler import RegistryHandlerBaseConfig
+from aiq.utils.data_models.schema_validator import validate_yaml
+
+logger = logging.getLogger(__name__)
+
+
+async def remove_artifact(registry_handler_config: RegistryHandlerBaseConfig, packages: list[dict[str, str]]) -> None:
+
+    from aiq.cli.type_registry import GlobalTypeRegistry
+    from aiq.registry_handlers.schemas.package import PackageNameVersionList
+
+    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))
+
+        try:
+            package_name_list = PackageNameVersionList(**{"packages": packages})
+        except Exception as e:
+            logger.exception("Invalid package format: '%s'", e, exc_info=True)
+
+        await stack.enter_async_context(registry_handler.remove(packages=package_name_list))
+
+
+@click.group(name=__name__,
+             invoke_without_command=True,
+             help=("Remove AgentIQ artifact from a remote registry by name and version."))
+@click.argument("packages", type=str)
+@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 YAML file to override the channel settings."),
+)
+@click.option(
+    "-c",
+    "--channel",
+    type=str,
+    required=True,
+    help=("The remote registry channel that will remove the AgentIQ artifact."),
+)
+def remove(channel: str, config_file: str, packages: str) -> None:
+    """
+    Remove AgentIQ artifacts from a remote registry.
+    """
+
+    from aiq.settings.global_settings import GlobalSettings
+
+    # Extract package name and version
+    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)
+
+    settings = GlobalSettings().get()
+
+    if (config_file is not None):
+        settings = settings.override_settings(config_file)
+
+    try:
+        remove_channel_config = settings.channels.get(channel)
+
+        if (remove_channel_config is 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)
+        return
+
+    asyncio.run(remove_artifact(registry_handler_config=remove_channel_config, packages=packages_versions))

aiq/cli/commands/registry/search.py

@@ -0,0 +1,155 @@
+# 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 AgentIQ 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 AgentIQ 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 AgentIQ 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))