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

aiq/registry_handlers/local/local_handler.py

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

@@ -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
-
-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

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

@@ -1,198 +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 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)