aiqtoolkit 1.1.0a20250503__py3-none-any.whl

aiq/registry_handlers/local/local_handler.py

@@ -0,0 +1,176 @@
+# 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.registry_handlers.package_utils import build_package_metadata
+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 PullRequestPackages
+from aiq.registry_handlers.schemas.pull import PullResponse
+from aiq.registry_handlers.schemas.remove import RemoveResponse
+from aiq.registry_handlers.schemas.search import SearchFields
+from aiq.registry_handlers.schemas.search import SearchQuery
+from aiq.registry_handlers.schemas.search import SearchResponse
+from aiq.registry_handlers.schemas.status import ActionEnum
+from aiq.registry_handlers.schemas.status import StatusEnum
+
+logger = logging.getLogger(__name__)
+
+
+class LocalRegistryHandler(AbstractRegistryHandler):
+    """A registry handler for interactions with a local Python environment."""
+
+    search_fields: list[SearchFields] = [field for field in SearchFields if field != SearchFields.ALL]
+
+    @asynccontextmanager
+    async def publish(self, artifact: AIQArtifact) -> AsyncGenerator[PublishResponse]:
+        """Publishes an AIQ Toolkit artifact to a local registry.
+
+        Args:
+            artifact (AIQArtifact): An artifact that contain AIQ Toolkit plugin wheel and it's corrosponding discovery
+            metadata.
+
+        Yields:
+            Iterator[AsyncGenerator[PublishResponse]]: A response message that includes a completion status message.
+        """
+
+        try:
+            validated_remove_response = RemoveResponse(status={
+                "status": StatusEnum.ERROR, "message": "Local publish not supported.", "action": ActionEnum.PUBLISH
+            })
+            yield validated_remove_response
+        finally:
+            logger.warning(validated_remove_response.status.message)
+
+    @asynccontextmanager
+    async def pull(self, packages: PullRequestPackages) -> AsyncGenerator[PullResponse]:
+        """Download and install AIQ Toolkit artifacts from a local 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.
+        """
+
+        try:
+            validated_remove_response = RemoveResponse(status={
+                "status": StatusEnum.ERROR, "message": "Local pull not supported.", "action": ActionEnum.PULL
+            })
+
+            yield validated_remove_response
+        finally:
+            logger.warning(validated_remove_response.status.message)
+
+    @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.
+        """
+
+        try:
+            results_dict = build_package_metadata(wheel_data=None)
+            component_results = []
+            query_component_types = [component_type.value for component_type in query.component_types]
+            for component_type, components in results_dict.items():
+                if component_type in query_component_types:
+                    component_results.extend(components)
+
+            if (SearchFields.ALL in query.fields):
+                query.fields = self.search_fields
+
+            matched_results = []
+            for component_result in component_results:
+                for search_field in query.fields:
+                    if (query.query in component_result.get(search_field.value, "")):
+                        matched_results.append(component_result)
+                        break
+
+            if query.top_k > 0:
+                top_k = query.top_k
+            else:
+                top_k = len(matched_results)
+
+            validated_search_response = SearchResponse(results=matched_results[:top_k],
+                                                       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}"
+            validated_search_response = SearchResponse(params=query,
+                                                       status={
+                                                           "status": StatusEnum.SUCCESS,
+                                                           "message": msg,
+                                                           "action": ActionEnum.SEARCH
+                                                       })
+            logger.exception(validated_search_response.status.message, exc_info=True)
+
+            yield validated_search_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    @asynccontextmanager
+    async def remove(self, packages: PackageNameVersionList) -> AsyncGenerator[RemoveResponse]:
+        """Uninstall packages from the local Python environment.
+
+        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:
+            for package_name in packages.packages:
+                result = subprocess.run(["uv", "pip", "uninstall", package_name.name], check=True)
+                result.check_returncode()
+
+            validated_remove_response = RemoveResponse(status={
+                "status": StatusEnum.SUCCESS, "message": "", "action": ActionEnum.REMOVE
+            })  # type: ignore
+
+            yield validated_remove_response
+
+        except Exception as e:
+            msg = f"Error uninstalling artifacts: {e}"
+            validated_remove_response = RemoveResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
+            })  # type: ignore
+            logger.exception(validated_remove_response.status.message, exc_info=True)
+
+            yield validated_remove_response
+
+        finally:
+            logger.info("Execution complete.")

aiq/registry_handlers/local/register_local.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
+
+from aiq.cli.register_workflow import register_registry_handler
+from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class LocalRegistryHandlerConfig(RegistryHandlerBaseConfig, name="local"):
+    """Interact with the local AIQ Toolkit environment to search and uninstall AIQ Toolkit components."""
+
+    pass
+
+
+@register_registry_handler(config_type=LocalRegistryHandlerConfig)
+async def local_registry_handler(config: LocalRegistryHandlerConfig):
+
+    from aiq.registry_handlers.local.local_handler import LocalRegistryHandler
+
+    registry_handler = LocalRegistryHandler()
+
+    yield registry_handler

aiq/registry_handlers/metadata_factory.py

@@ -0,0 +1,60 @@
+# 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
+
+from aiq.data_models.component import AIQComponentEnum
+from aiq.data_models.discovery_metadata import DiscoveryMetadata
+from aiq.data_models.discovery_metadata import DiscoveryStatusEnum
+from aiq.registry_handlers.schemas.package import WheelData
+
+logger = logging.getLogger(__name__)
+
+
+class ComponentDiscoveryMetadata:
+
+    def __init__(self, component_type: AIQComponentEnum, wheel_data: WheelData | None = None):
+        self._component_type = component_type
+        self._metadata_items: list[dict | DiscoveryMetadata] = []
+        self._wheel_data: WheelData = wheel_data
+
+    def load_metadata(self):
+
+        from aiq.cli.type_registry import GlobalTypeRegistry
+
+        registry = GlobalTypeRegistry.get()
+
+        for _, registered_component_info in registry.get_infos_by_type(component_type=self._component_type).items():
+
+            if ((registered_component_info.discovery_metadata.status == DiscoveryStatusEnum.SUCCESS) and
+                ((self._wheel_data is None) or
+                 (registered_component_info.discovery_metadata.package in self._wheel_data.union_dependencies))):
+
+                if ((self._wheel_data is not None)
+                        and (registered_component_info.discovery_metadata.package == self._wheel_data.package_name)):
+                    discovery_metadata_copy = registered_component_info.discovery_metadata.model_copy(deep=True)
+                    discovery_metadata_copy.version = self._wheel_data.whl_version
+                    self._metadata_items.append(discovery_metadata_copy.model_dump())
+                    continue
+
+                self._metadata_items.append(registered_component_info.discovery_metadata.model_dump())
+
+    def get_metadata_items(self) -> list[dict | DiscoveryMetadata]:
+        return self._metadata_items
+
+    @staticmethod
+    def from_package_component_type(component_type: AIQComponentEnum,
+                                    wheel_data: WheelData | None = None) -> "ComponentDiscoveryMetadata":
+        return ComponentDiscoveryMetadata(component_type=component_type, wheel_data=wheel_data)

aiq/registry_handlers/package_utils.py

@@ -0,0 +1,198 @@
+# 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 base64
+import importlib.metadata
+import logging
+import os
+import subprocess
+from functools import lru_cache
+
+from aiq.data_models.component import AIQComponentEnum
+from aiq.data_models.discovery_metadata import DiscoveryMetadata
+from aiq.registry_handlers.schemas.package import WheelData
+from aiq.registry_handlers.schemas.publish import AIQArtifact
+from aiq.runtime.loader import PluginTypes
+from aiq.runtime.loader import discover_entrypoints
+
+# pylint: disable=redefined-outer-name
+logger = logging.getLogger(__name__)
+
+
+@lru_cache
+def get_module_name_from_distribution(distro_name: str) -> str | None:
+    """Return the first top-level module name for a given distribution name."""
+    if not distro_name:
+        return None
+
+    try:
+        # Read 'top_level.txt' which contains the module(s) provided by the package
+        dist = importlib.metadata.distribution(distro_name)
+        # will reading a file set of vun scan?
+        top_level = dist.read_text('top_level.txt')
+
+        if top_level:
+            module_names = top_level.strip().split()
+            # return firs module name
+            return module_names[0]
+    except importlib.metadata.PackageNotFoundError:
+        # Distribution not found
+        return None
+    except FileNotFoundError:
+        # 'top_level.txt' might be missing
+        return None
+
+    return None
+
+
+def build_wheel(package_root: str) -> WheelData:
+    """Builds a Python .whl for the specified package and saves to disk, sets self._whl_path, and returned as bytes.
+
+    Args:
+        package_root (str): Path to the local package repository.
+
+    Returns:
+        WheelData: Data model containing a built python wheel and its corresponding metadata.
+    """
+
+    import importlib.util
+    import re
+    import tomllib
+
+    from pkginfo import Wheel
+
+    pyproject_toml_path = os.path.join(package_root, "pyproject.toml")
+
+    if not os.path.exists(pyproject_toml_path):
+        raise ValueError("Invalid package path, does not contain a pyproject.toml file.")
+
+    with open(pyproject_toml_path, "rb") as f:
+        data = tomllib.load(f)
+
+    toml_project: dict = data.get("project", {})
+    toml_project_name = toml_project.get("name", None)
+
+    assert toml_project_name is not None, f"Package name '{toml_project_name}' not found in pyproject.toml"
+    # replace "aiqtoolkit" substring with "aiq" to get the import name
+    module_name = get_module_name_from_distribution(toml_project_name)
+    assert module_name is not None, f"No modules found for package name '{toml_project_name}'"
+
+    assert importlib.util.find_spec(module_name) is not None, (f"Package {module_name} not "
+                                                               "installed, cannot discover components.")
+
+    toml_packages = set(i for i in data.get("project", {}).get("entry-points", {}).get("aiq.plugins", {}))
+    toml_dependencies = set(
+        re.search(r"[a-zA-Z][a-zA-Z\d_-]*", package_name).group(0)
+        for package_name in toml_project.get("dependencies", []))
+
+    union_dependencies = toml_dependencies.union(toml_packages)
+    union_dependencies.add(toml_project_name)
+
+    working_dir = os.getcwd()
+    os.chdir(package_root)
+
+    result = subprocess.run(["uv", "build", "--wheel"], check=True)
+    result.check_returncode()
+
+    whl_file = sorted(os.listdir("dist"), reverse=True)[0]
+    whl_file_path = os.path.join("dist", whl_file)
+
+    with open(whl_file_path, "rb") as whl:
+        whl_bytes = whl.read()
+        whl_base64 = base64.b64encode(whl_bytes).decode("utf-8")
+
+    whl_path = os.path.join(os.getcwd(), whl_file_path)
+
+    os.chdir(working_dir)
+
+    whl_version = Wheel(whl_path).version
+
+    return WheelData(
+        package_root=package_root,
+        package_name=module_name,  # should it be module name or distro name here
+        toml_project=toml_project,
+        toml_dependencies=toml_dependencies,
+        toml_aiq_packages=toml_packages,
+        union_dependencies=union_dependencies,
+        whl_path=whl_path,
+        whl_base64=whl_base64,
+        whl_version=whl_version)
+
+
+def build_package_metadata(wheel_data: WheelData | None) -> dict[AIQComponentEnum, list[dict | DiscoveryMetadata]]:
+    """Loads discovery metadata for all registered AIQ Toolkit components included in this Python package.
+
+    Args:
+        wheel_data (WheelData): Data model containing a built python wheel and its corresponding metadata.
+
+    Returns:
+        dict[AIQComponentEnum, list[typing.Union[dict, DiscoveryMetadata]]]: List containing each components discovery
+        metadata.
+    """
+
+    from aiq.cli.type_registry import GlobalTypeRegistry
+    from aiq.registry_handlers.metadata_factory import ComponentDiscoveryMetadata
+    from aiq.runtime.loader import discover_and_register_plugins
+
+    discover_and_register_plugins(PluginTypes.ALL)
+
+    registry = GlobalTypeRegistry.get()
+
+    aiq_plugins = discover_entrypoints(PluginTypes.ALL)
+
+    if (wheel_data is not None):
+        registry.register_package(package_name=wheel_data.package_name, package_version=wheel_data.whl_version)
+        for entry_point in aiq_plugins:
+            package_name = entry_point.module.split('.')[0]
+            if (package_name == wheel_data.package_name):
+                continue
+            if (package_name in wheel_data.union_dependencies):
+                registry.register_package(package_name=package_name)
+
+    else:
+        for entry_point in aiq_plugins:
+            package_name = entry_point.module.split('.')[0]
+            registry.register_package(package_name=package_name)
+
+    discovery_metadata = {}
+    for component_type in AIQComponentEnum:
+
+        if (component_type == AIQComponentEnum.UNDEFINED):
+            continue
+        component_metadata = ComponentDiscoveryMetadata.from_package_component_type(wheel_data=wheel_data,
+                                                                                    component_type=component_type)
+        component_metadata.load_metadata()
+        discovery_metadata[component_type] = component_metadata.get_metadata_items()
+
+    return discovery_metadata
+
+
+def build_aiq_artifact(package_root: str) -> AIQArtifact:
+    """Builds a complete AIQ Toolkit Artifact that can be published for discovery and reuse.
+
+    Args:
+        package_root (str): Path to root of python package
+
+    Returns:
+        AIQArtifact: An publishabla AIQArtifact containing package wheel and discovery metadata.
+    """
+
+    from aiq.registry_handlers.schemas.publish import BuiltAIQArtifact
+
+    wheel_data = build_wheel(package_root=package_root)
+    metadata = build_package_metadata(wheel_data=wheel_data)
+    built_artifact = BuiltAIQArtifact(whl=wheel_data.whl_base64, metadata=metadata)
+
+    return AIQArtifact(artifact=built_artifact, whl_path=wheel_data.whl_path)