nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/registry_handlers/pypi/pypi_handler.py

@@ -0,0 +1,251 @@
+# 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 subprocess
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from aiq.data_models.component import AIQComponentEnum
+from aiq.registry_handlers.registry_handler_base import AbstractRegistryHandler
+from aiq.registry_handlers.schemas.package import PackageNameVersionList
+from aiq.registry_handlers.schemas.publish import AIQArtifact
+from aiq.registry_handlers.schemas.publish import PublishResponse
+from aiq.registry_handlers.schemas.pull import PackageNameVersion
+from aiq.registry_handlers.schemas.pull import PullRequestPackages
+from aiq.registry_handlers.schemas.pull import PullResponse
+from aiq.registry_handlers.schemas.remove import RemoveResponse
+from aiq.registry_handlers.schemas.search import SearchQuery
+from aiq.registry_handlers.schemas.search import SearchResponse
+from aiq.registry_handlers.schemas.search import SearchResponseItem
+from aiq.registry_handlers.schemas.status import ActionEnum
+from aiq.registry_handlers.schemas.status import StatusEnum
+
+logger = logging.getLogger(__name__)
+
+
+class PypiRegistryHandler(AbstractRegistryHandler):
+    """
+    A registry handler for interactions with a remote PyPI registry.
+
+    Built interfacing with this private PyPI server:
+    https://github.com/pypiserver/pypiserver
+    """
+
+    def __init__(  # pylint: disable=R0917
+            self,
+            endpoint: str,
+            token: str | None = None,
+            publish_route: str = "",
+            pull_route: str = "",
+            search_route: str = ""):
+        super().__init__()
+        self._endpoint = endpoint.rstrip("/")
+        self._token = token
+        self._publish_route = publish_route.strip("/")
+        self._pull_route = pull_route.strip("/")
+        self._search_route = search_route.strip("/")
+
+    @asynccontextmanager
+    async def publish(self, artifact: AIQArtifact) -> AsyncGenerator[PublishResponse]:
+        """Publishes an AIQ Toolkit artifact to a PyPI remote registry.
+
+        Args:
+            artifact (AIQArtifact): An artifact that contain AIQ Toolkit plugin wheel and it's corrosponding discovery
+            metadata.
+
+        Yields:
+            Iterator[AsyncGenerator[PublishResponse, None]]: A response message that includes a completion status
+            message.
+        """
+
+        try:
+            result = self._upload_to_pypi(wheel_path=artifact.whl_path)
+            result.check_returncode()
+
+            validated_publish_response = PublishResponse(status={
+                "status": StatusEnum.SUCCESS, "message": "", "action": ActionEnum.PUBLISH
+            })
+
+            yield validated_publish_response
+
+        except Exception as e:
+            msg = f"Error publishing package: {e}"
+            validated_publish_response = PublishResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PUBLISH
+            })
+            logger.exception(validated_publish_response.status.message, exc_info=True)
+
+            yield validated_publish_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    def _upload_to_pypi(self, wheel_path: str) -> None:
+
+        return subprocess.run(
+            ["twine", "upload", "--repository-url", f"{self._endpoint}/{self._publish_route}", f"{wheel_path}"],
+            check=True)
+
+    @asynccontextmanager
+    async def pull(self, packages: PullRequestPackages) -> AsyncGenerator[PullResponse]:
+        """Download and install AIQ Toolkit artifacts from a remote PyPI remote registry.
+
+        Args:
+            packages (PullRequestPackages): Parameters used to pull the AIQ Toolkit artifact.
+
+        Yields:
+            Iterator[AsyncGenerator[PullResponse, None]]: A response message that includes a the pulled packages and a
+                completion status message.
+        """
+
+        try:
+            versioned_packages = []
+            for package in packages.packages:
+
+                if isinstance(package, PackageNameVersion):
+                    if (hasattr(package, "version") and package.version is not None):
+                        versioned_package = f"{package.name}=={package.version}"
+                    else:
+                        versioned_package = package.name
+                    versioned_packages.append(versioned_package)
+                else:
+                    versioned_packages.append(package.whl_path)
+
+            versioned_packages_str = " ".join(versioned_packages)
+
+            result = subprocess.run(
+                [
+                    "uv",
+                    "pip",
+                    "install",
+                    "--prerelease=allow",
+                    "--index-url",
+                    f"{self._endpoint}/{self._pull_route}/",
+                    versioned_packages_str
+                ],  # pylint: disable=W0631
+                check=True)
+
+            result.check_returncode()
+
+            validated_pull_response = PullResponse(status={
+                "status": StatusEnum.SUCCESS, "message": "", "action": ActionEnum.PULL
+            })
+
+            yield validated_pull_response
+
+        except Exception as e:
+            msg = f"Error pulling packages: {e}"
+            validated_pull_response = PullResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PULL
+            })
+            logger.exception(validated_pull_response.status.message, exc_info=True)
+
+            yield validated_pull_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    @asynccontextmanager
+    async def search(self, query: SearchQuery) -> AsyncGenerator[SearchResponse]:
+        """Searches a remote PyPI registry for relevant AIQ Toolkit components.
+
+        Args:
+            query (SearchQuery): Parameters of the search to be performed.
+
+        Yields:
+            Iterator[AsyncGenerator[SearchResponse]]: A response message that includes search
+                parameters and a completion status message.
+        """
+
+        try:
+            completed_process = subprocess.run(
+                ["pip", "search", "--index", f"{self._endpoint}", query.query],  # pylint: disable=W0631
+                text=True,
+                capture_output=True,
+                check=True)
+            search_response_list = []
+            search_results = completed_process.stdout
+            package_results = search_results.split("\n")
+
+            for package_result in package_results:
+
+                # Filter out empty and nested values
+                if ((package_result == "") or (package_result[0] == " ")):
+                    continue
+
+                package_split = package_result.split(" ")
+                package = package_split[0]
+                version = package_split[1][1:-1]
+
+                search_resp_item = SearchResponseItem(package=package,
+                                                      version=version,
+                                                      component_type=AIQComponentEnum.PACKAGE,
+                                                      component_name=package,
+                                                      description="",
+                                                      developer_notes="")
+
+                if (search_resp_item not in search_response_list):
+                    search_response_list.append(search_resp_item)
+
+                    if (len(search_response_list) > query.top_k):
+                        break
+
+            validated_search_response = SearchResponse(results=search_response_list,
+                                                       params=query,
+                                                       status={
+                                                           "status": StatusEnum.SUCCESS,
+                                                           "message": "",
+                                                           "action": ActionEnum.SEARCH
+                                                       })
+
+            yield validated_search_response
+
+        except Exception as e:
+            msg = f"Error searching for artifacts: {e}"
+            logger.exception(msg, exc_info=True)
+            validated_search_response = SearchResponse(params=query,
+                                                       status={
+                                                           "status": StatusEnum.ERROR,
+                                                           "message": msg,
+                                                           "action": ActionEnum.SEARCH
+                                                       })
+
+            yield validated_search_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    @asynccontextmanager
+    async def remove(self, packages: PackageNameVersionList) -> AsyncGenerator[SearchResponse]:
+        """Removes packages from a remote registry.
+
+        Args:
+            packages (PackageNameVersionList): The list of packages to remove.
+
+        Yields:
+            Iterator[AsyncGenerator[RemoveResponse]]: A response message that includes the packages and a
+                completion status message.
+        """
+
+        try:
+            msg = "PyPI remove not supported."
+            validated_remove_response = RemoveResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
+            })
+
+            yield validated_remove_response
+        finally:
+            logger.warning(validated_remove_response.status.message)

aiq/registry_handlers/pypi/register_pypi.py

@@ -0,0 +1,40 @@
+# 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.
+
+from pydantic import Field
+
+from aiq.cli.register_workflow import register_registry_handler
+from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
+
+
+class PypiRegistryHandlerConfig(RegistryHandlerBaseConfig, name="pypi"):
+    """Registry handler for interacting with a remote PyPI registry index."""
+
+    endpoint: str = Field(description="A string representing the remote endpoint.")
+    token: str | None = Field(default=None,
+                              description="The authentication token to use when interacting with the registry.")
+    publish_route: str = Field(description="The route to the AIQ Toolkit publish service.")
+    pull_route: str = Field(description="The route to the AIQ Toolkit pull service.")
+    search_route: str = Field(default="simple", description="The route to the AIQ Toolkit search service.")
+
+
+@register_registry_handler(config_type=PypiRegistryHandlerConfig)
+async def pypi_publish_registry_handler(config: PypiRegistryHandlerConfig):
+
+    from aiq.registry_handlers.pypi.pypi_handler import PypiRegistryHandler
+
+    registry_handler = PypiRegistryHandler(endpoint=config.endpoint, token=config.token)
+
+    yield registry_handler

aiq/registry_handlers/register.py

@@ -0,0 +1,21 @@
+# 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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+
+from .local import register_local  # pylint: disable=E0611
+from .pypi import register_pypi  # pylint: disable=E0611
+from .rest import register_rest  # pylint: disable=E0611

aiq/registry_handlers/registry_handler_base.py

@@ -0,0 +1,157 @@
+# 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.
+
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from enum import Enum
+
+from aiq.data_models.component import AIQComponentEnum
+from aiq.data_models.discovery_metadata import DiscoveryMetadata
+from aiq.registry_handlers.schemas.package import PackageNameVersionList
+from aiq.registry_handlers.schemas.publish import AIQArtifact
+from aiq.registry_handlers.schemas.publish import PublishResponse
+from aiq.registry_handlers.schemas.pull import PullRequestPackages
+from aiq.registry_handlers.schemas.pull import PullResponse
+from aiq.registry_handlers.schemas.remove import RemoveResponse
+from aiq.registry_handlers.schemas.search import SearchQuery
+from aiq.registry_handlers.schemas.search import SearchResponse
+from aiq.registry_handlers.schemas.search import VisualizeFields
+
+
+class AbstractRegistryHandler(ABC):
+    """Base class outlining the interfaces for remote AIQ Toolkit registry interactions."""
+
+    def __init__(self):
+        self._discovery_metadata: dict[AIQComponentEnum, list[dict | DiscoveryMetadata]] = {}
+        self._aiq_artifact: AIQArtifact | None = None
+        self._whl_bytes: bytes
+        self._whl_path: str
+        self._whl_base64: str
+
+    @abstractmethod
+    @asynccontextmanager
+    async def publish(self, artifact: AIQArtifact) -> AsyncGenerator[PublishResponse]:
+        """Publishes an AIQ Toolkit artifact to a remote registry.
+
+        Args:
+            artifact (AIQArtifact): An artifact that contain AIQ Toolkit plugin wheel and it's corrosponding discovery
+            metadata.
+
+        Yields:
+            Iterator[AsyncGenerator[PublishResponse, None]]: A response message that includes a completion status
+            message.
+        """
+
+        pass
+
+    @abstractmethod
+    @asynccontextmanager
+    async def pull(self, packages: PullRequestPackages) -> AsyncGenerator[PullResponse]:
+        """Download and install AIQ Toolkit artifacts from a remote registry.
+
+        Args:
+            packages (PullRequestPackages): Parameters used to pull the AIQ Toolkit artifact.
+
+        Yields:
+            Iterator[AsyncGenerator[PullResponse]]: A response message that includes a the pulled packages and a
+                completion status message.
+        """
+
+        pass
+
+    @abstractmethod
+    @asynccontextmanager
+    async def search(self, query: SearchQuery) -> AsyncGenerator[SearchResponse]:
+        """Searches the local aiq registry for relevant AIQ Toolkit components.
+
+        Args:
+            query (SearchQuery): Parameters of the search to be performed.
+
+        Yields:
+            Iterator[AsyncGenerator[SearchResponse]]: A response message that includes search
+                parameters and a completion status message.
+        """
+
+        pass
+
+    @abstractmethod
+    @asynccontextmanager
+    async def remove(self, packages: PackageNameVersionList) -> AsyncGenerator[RemoveResponse]:
+        """Removes packages from a remote registry.
+
+        Args:
+            packages (PackageNameVersionList): The list of packages to remove.
+
+        Yields:
+            Iterator[AsyncGenerator[RemoveResponse]]: A response message that includes the packages and a
+                completion status message.
+        """
+
+        pass
+
+    @staticmethod
+    def visualize_search_results(search_response: SearchResponse, pager: bool = True) -> None:
+        """Visualze search results in a system terminal.
+
+        Args:
+            search_response (SearchResponse): A response message that includes search parameters and a completion status
+            message.
+
+            pager (bool, optional): Include an pagable terminal interface for large search results. Defaults to False.
+        """
+
+        from rich.console import Console
+        from rich.table import Table
+        from rich.text import Text
+
+        table = Table(title="AIQ Toolkit Search Results", padding=(0, 1), show_lines=True)
+        for column in VisualizeFields:
+            table.add_column(column.value)
+
+        for result in search_response.results:
+            row = []
+            for column in VisualizeFields:
+                value = getattr(result, column.value)
+                if isinstance(value, Enum):
+                    value = value.value
+                text = Text(value, overflow="fold")
+                row.append(text)
+            table.add_row(*row, style='bright_green')
+
+        console = Console()
+
+        if (pager):
+            with console.pager():
+                console.print(table)
+        else:
+            console.print(table)
+
+    @staticmethod
+    def save_search_results(search_response: SearchResponse, save_path: str) -> None:
+        """Save search results to a local json file.
+
+        Args:
+            search_response (SearchResponse): A response message that includes search parameters and a completion status
+            message.
+
+            save_path (str): The path to save the json search results.
+        """
+
+        search_response_str = search_response.model_dump_json(indent=4)
+
+        with open(save_path, "w", encoding="utf-8") as f:
+            f.write(search_response_str)

aiq/registry_handlers/rest/register_rest.py

@@ -0,0 +1,56 @@
+# 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 os
+
+from pydantic import Field
+
+from aiq.cli.register_workflow import register_registry_handler
+from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
+
+
+class RestRegistryHandlerConfig(RegistryHandlerBaseConfig, name="rest"):
+    """Registry handler for interacting with a remote REST registry."""
+
+    endpoint: str = Field(description="A string representing the remote endpoint.")
+    token: str | None = Field(default=None,
+                              description="The authentication token to use when interacting with the registry.")
+    publish_route: str = Field(default="", description="The route to the AIQ Toolkit publish service.")
+    pull_route: str = Field(default="", description="The route to the AIQ Toolkit pull service.")
+    search_route: str = Field(default="", description="The route to the AIQ Toolkit search service")
+    remove_route: str = Field(default="", description="The route to the AIQ Toolkit remove service")
+
+
+@register_registry_handler(config_type=RestRegistryHandlerConfig)
+async def rest_search_handler(config: RestRegistryHandlerConfig):
+
+    from aiq.registry_handlers.rest.rest_handler import RestRegistryHandler
+
+    if (config.token is None):
+        registry_token = os.getenv("REGISTRY_TOKEN")
+
+        if (registry_token is None):
+            raise ValueError("Please supply registry token.")
+    else:
+        registry_token = config.token
+
+    registry_handler = RestRegistryHandler(token=registry_token,
+                                           endpoint=config.endpoint,
+                                           publish_route=config.publish_route,
+                                           pull_route=config.pull_route,
+                                           search_route=config.search_route,
+                                           remove_route=config.remove_route)
+
+    yield registry_handler