aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/registry_handlers/rest/rest_handler.py

@@ -1,237 +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 logging
-import os
-import shutil
-import subprocess
-from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
-
-import httpx
-
-from aiq.registry_handlers.registry_handler_base import AbstractRegistryHandler
-from aiq.registry_handlers.schemas.headers import RequestHeaders
-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.status import ActionEnum
-from aiq.registry_handlers.schemas.status import StatusEnum
-
-logger = logging.getLogger(__name__)
-
-
-class RestRegistryHandler(AbstractRegistryHandler):
-    """A registry handler for interactions with a remote REST registry."""
-
-    def __init__(  # pylint: disable=R0917
-            self,
-            endpoint: str,
-            token: str,
-            timeout: int = 30,
-            publish_route: str = "",
-            pull_route: str = "",
-            search_route: str = "",
-            remove_route: str = ""):
-        super().__init__()
-        self._endpoint = endpoint.rstrip("/")
-        self._timeout = timeout
-        self._publish_route = publish_route.strip("/")
-        self._pull_route = pull_route.strip("/")
-        self._search_route = search_route.strip("/")
-        self._remove_route = remove_route.strip("/")
-        self._headers = RequestHeaders(Authorization=f"Bearer: {token}").model_dump(by_alias=True)
-
-    @asynccontextmanager
-    async def publish(self, artifact: AIQArtifact) -> AsyncGenerator[PublishResponse]:
-        """Publishes an AIQ Toolkit artifact to a remote REST 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:
-
-            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
-                response = await client.post(f"{self._endpoint}/{self._publish_route}",
-                                             content=artifact.artifact.model_dump_json())
-                response.raise_for_status()
-                response_dict = response.json()
-
-            validated_publish_response = PublishResponse(**response_dict)
-
-            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.")
-
-    @asynccontextmanager
-    async def pull(self, packages: PullRequestPackages) -> AsyncGenerator[PullResponse]:
-        """Download and install AIQ Toolkit artifacts from a remote REST 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.
-        """
-
-        tmp_dir = ".tmp-aiq-pull"
-
-        try:
-            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
-                packages_json = packages.model_dump_json()
-                response = await client.post(f"{self._endpoint}/{self._pull_route}", content=packages_json)
-                response.raise_for_status()
-
-            response_dict = response.json()
-            validated_pull_response = PullResponse(**response_dict)
-
-            if (validated_pull_response.status.status == StatusEnum.ERROR):
-                logger.error(validated_pull_response.status.message)
-                raise ValueError(f"Server error: {validated_pull_response.status.message}")
-
-            if (not os.path.exists(tmp_dir)):
-                os.mkdir(tmp_dir)
-
-            whl_paths = []
-
-            for package in validated_pull_response.packages:
-                whl_bytes = base64.b64decode(package.whl)
-                whl_path = os.path.join(tmp_dir, package.whl_name)
-
-                with open(whl_path, "wb") as f:
-                    f.write(whl_bytes)
-
-                whl_paths.append(whl_path)
-
-            cmd = ["uv", "pip", "install"]
-            cmd.extend(whl_paths)
-            result = subprocess.run(cmd, check=True)
-            result.check_returncode()
-
-            if (os.path.exists(tmp_dir)):
-                shutil.rmtree(tmp_dir)
-
-            yield validated_pull_response
-
-        except Exception as e:
-            msg = f"Error occured when installing packages: {e}"
-            logger.error(msg)
-            if (os.path.exists(tmp_dir)):
-                shutil.rmtree(tmp_dir)
-
-            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 REST 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:
-            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
-                query_json = query.model_dump_json()
-                response = await client.post(url=f"{self._endpoint}/{self._search_route}", content=query_json)
-                response.raise_for_status()
-
-            response_dict = response.json()
-            validated_search_response = SearchResponse(**response_dict)
-
-            yield validated_search_response
-
-        except Exception as e:
-            msg = f"Error searching for artifacts: {e}"
-            validated_search_response = SearchResponse(params=query,
-                                                       status={
-                                                           "status": StatusEnum.ERROR,
-                                                           "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]:
-        """Removes packages from a remote REST 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:
-            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
-                remove_json = packages.model_dump_json()
-                response = await client.post(url=f"{self._endpoint}/{self._remove_route}", content=remove_json)
-                response.raise_for_status()
-
-            response_dict = response.json()
-            validated_remove_response = RemoveResponse(**response_dict)
-
-            yield validated_remove_response
-
-        except Exception as e:
-            msg = f"Error removing artifacts: {e}"
-            validated_remove_response = RemoveResponse(status={
-                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
-            })
-            logger.exception(validated_remove_response.status.message, exc_info=True)
-
-            yield validated_remove_response
-
-        finally:
-            logger.info("Execution complete.")

aiq/registry_handlers/schemas/headers.py

@@ -1,42 +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 pydantic import BaseModel
-from pydantic import Field
-
-logger = logging.getLogger(__name__)
-
-
-class RequestHeaders(BaseModel):
-    """Represents a data model for REST registry handler request headers.
-
-    Args:
-        accept (str): Specifies the media types the client can accept. Defaults to 'application/json'
-        content_type (str): Describes the format of the request body data. Defaults to 'application/json'
-        authorization (str): Contains authentication credentials for accessing a protected resource.
-    """
-
-    accept: str = Field("application/json", alias="Accept")
-    content_type: str = Field("application/json", alias="Content-Type")
-    authorization: str = Field(..., alias="Authorization")
-
-
-class ResponseHeaders(BaseModel):
-    """Placehoder data model for REST registry handler resopnse headers.
-    """
-
-    pass

aiq/registry_handlers/schemas/package.py

@@ -1,68 +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 pydantic import BaseModel
-
-logger = logging.getLogger(__name__)
-
-
-class WheelData(BaseModel):
-    """Data model containing a built python wheel and its corresponding metadata.
-
-    Args:
-        package_root (str): The path to the package root directory containing the pyproject.toml file.
-        package_name (str): The name of the python package.
-        toml_project (dict): A dictionary containing data about the python project.
-        toml_dependencies (set): The list of dependencies provided in the pyproject.toml file.
-        toml_aiq_packages (set): The AIQ Toolkit plugins listed in the pyproject.toml.
-        union_dependencies (set): The union of toml_dependencies and toml_aiq_packages.
-        whl_path (str): The path to the package wheel file.
-        whl_base64 (str): Base64 encoded string of the wheel file.
-        whl_version (str): The version representing the wheel file.
-    """
-
-    package_root: str
-    package_name: str
-    toml_project: dict
-    toml_dependencies: set
-    toml_aiq_packages: set
-    union_dependencies: set
-    whl_path: str
-    whl_base64: str
-    whl_version: str
-
-
-class PackageNameVersion(BaseModel):
-    """Represents a data model containing a package name and version.
-
-    Args:
-        name (str): Package name, excluding the version.
-        version (str | None): The package version, excluding the name. Defaults to None.
-    """
-
-    name: str
-    version: str | None = None
-
-
-class PackageNameVersionList(BaseModel):
-    """Represents a data model containing a list of `PackageNameVersion` packages.
-
-    Args:
-        packages (list[PackageNameVersion]): A list of `PackageNameVersion` models. Defaults to `[]`
-    """
-
-    packages: list[PackageNameVersion] = []

aiq/registry_handlers/schemas/publish.py

@@ -1,63 +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 pydantic import BaseModel
-
-from aiq.data_models.component import AIQComponentEnum
-from aiq.data_models.discovery_metadata import DiscoveryMetadata
-from aiq.registry_handlers.schemas.status import StatusMessage
-
-logger = logging.getLogger(__name__)
-
-
-class BuiltAIQArtifact(BaseModel):
-    """An AIQ Toolkit artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
-
-    Args:
-        whl (str): A base64 encoded string of an AIQ Toolkit package wheel (.whl).
-
-        metadata (dict[AIQComponentEnum, list[DiscoveryMetadata]]): Provides rich discover metadata for developers to
-        quickly find useful components.
-    """
-
-    whl: str
-    metadata: dict[AIQComponentEnum, list[DiscoveryMetadata]]
-
-
-class AIQArtifact(BaseModel):
-    """An AIQ Toolkit artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
-
-    Args:
-        artifact (BuildAIQArtifact): An AIQ Toolkit artifact including base64 encoded string of wheel package and
-        corrosponding discovery metadata.
-
-        whl_path (str): A local path to the built wheel package.
-    """
-
-    artifact: BuiltAIQArtifact | None = None
-    whl_path: str
-
-
-class PublishResponse(BaseModel):
-    """The expected response from a publish request denoting status information.
-
-    Args:
-        status (StatusMessage): Provides metadata describing the success or errors that occurred when
-            making a publish request.
-    """
-
-    status: StatusMessage

aiq/registry_handlers/schemas/pull.py

@@ -1,82 +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 pydantic import BaseModel
-
-from aiq.registry_handlers.schemas.package import PackageNameVersion
-from aiq.registry_handlers.schemas.status import StatusMessage
-
-logger = logging.getLogger(__name__)
-
-
-class PulledPackage(BaseModel):
-    """Represents a data model of a pulled package containing the package wheel and its name.
-
-    Args:
-        whl (str): Base64 encoded string of the AIQ Toolkit python package wheel (.whl).
-        whl_name (str): A string representing the wheel filename.
-    """
-
-    whl: str
-    whl_name: str
-
-
-class PullResponse(BaseModel):
-    """
-    Represents a data model of the expected respones from a AIQ Toolkit pull request, including detailed status
-    information.
-
-    Args:
-        packages (list[PulledPackage]): A list of pulled packages included in the pull request.
-        status (StatusMessage): Provides metadata describing the success or errors that occurred when making to pull in
-        a package.
-    """
-
-    packages: list[PulledPackage] = []
-    status: StatusMessage
-
-
-class PullPackageWhl(BaseModel):
-    """Local path to wheel (.whl) file.
-
-    Args:
-        whl_path (str): The local path the wheel (.whl) file.
-    """
-
-    whl_path: str
-
-
-class PullRequestPackage(BaseModel):
-    """Represents all data for a single package needed to download an install its components.
-
-    Args:
-        package (typing.Union[PackageNameVersion, PullPackageWhl]): Attributes of a single package necessary
-            to download and install its components.
-    """
-
-    package: PackageNameVersion | PullPackageWhl
-
-
-class PullRequestPackages(BaseModel):
-    """Represents a list of all packages th download and install in the local AIQ Toolkit environment.
-
-    Args:
-        packages (list[typing.Union[PackageNameVersion, PullPackageWhl]]): A list of packages that can be
-            downloaded and installed in the local AIQ Toolkit environment.
-    """
-
-    packages: list[PackageNameVersion | PullPackageWhl]

aiq/registry_handlers/schemas/remove.py

@@ -1,36 +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 pydantic import BaseModel
-
-from aiq.registry_handlers.schemas.package import PackageNameVersion
-from aiq.registry_handlers.schemas.status import StatusMessage
-
-logger = logging.getLogger(__name__)
-
-
-class RemoveResponse(BaseModel):
-    """Represents a data model for the expected response from a remove request, including packages and status metadata.
-
-    Args:
-        packages (list[PackageNameVersion]): A list of packages that are to be removed from a remote registry.
-        status (StatusMessage): Provides metadata describing the success or errors that occurred when making a remove
-        request.
-    """
-
-    packages: list[PackageNameVersion] = []
-    status: StatusMessage

aiq/registry_handlers/schemas/search.py

@@ -1,91 +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 enum import Enum
-
-from pydantic import BaseModel
-
-from aiq.data_models.component import AIQComponentEnum
-from aiq.registry_handlers.schemas.status import StatusMessage
-
-logger = logging.getLogger(__name__)
-
-
-class SearchFields(str, Enum):
-    ALL = "all"
-    PACKAGE = "package"
-    VERSION = "version"
-    COMPONENT_NAME = "component_name"
-    DESCRIPTION = "description"
-    DEVELOPER_NOTES = "developer_notes"
-
-
-class VisualizeFields(str, Enum):
-    PACKAGE = "package"
-    VERSION = "version"
-    COMPONENT_TYPE = "component_type"
-    COMPONENT_NAME = "component_name"
-    DESCRIPTION = "description"
-
-
-class SearchQuery(BaseModel):
-    """Represents the search criteria that will be used to discover useful AIQ Toolkit components.
-
-    Args:
-        query (str): A query string used to find useful AIQ Toolkit components.
-        fields (list[SearchFields]): The list of fields used when applying the query string.
-        component_types (list[AIQComponentEnum]): AIQ Toolkit components types to filter search results.
-        top_k (int): Specifies the number of search results to provide.
-    """
-
-    query: str = "*"
-    fields: list[SearchFields] = [SearchFields.ALL]
-    component_types: list[AIQComponentEnum]
-    top_k: int = 10
-
-
-class SearchResponseItem(BaseModel):
-    """Represents an individual item in the search response, including elements of it's discovery metadata.
-
-    Args:
-        package (str): The name of the AIQ Toolkit package that includes the component.
-        version (str): The version of the AIQ Toolkit package that includes the component.
-        component_type (AIQComponentEnum): Type of AIQ Toolkit component this item represents.
-        description (str): A description of this AIQ Toolkit component.
-        developer_notes (str): Additional details that would help a developer use this component.
-    """
-
-    package: str
-    version: str
-    component_type: AIQComponentEnum
-    component_name: str
-    description: str
-    developer_notes: str
-
-
-class SearchResponse(BaseModel):
-    """Represents a data model of the expected search response.
-
-    Args:
-        results (list[SearchResponseItem]): A list of results that matched the search criteria.
-        params (SearchQuery): The search criterial that produced these search results.
-        status (StatusMessage): Provides metadata describing the success or errors that occurred when making the search
-        request.
-    """
-
-    results: list[SearchResponseItem] = []
-    params: SearchQuery
-    status: StatusMessage

aiq/registry_handlers/schemas/status.py

@@ -1,47 +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 enum import Enum
-
-from pydantic import BaseModel
-
-logger = logging.getLogger(__name__)
-
-
-class ActionEnum(str, Enum):
-    PUBLISH = "publish"
-    PULL = "pull"
-    REMOVE = "remove"
-    SEARCH = "search"
-
-
-class StatusEnum(str, Enum):
-    SUCCESS = "success"
-    ERROR = "error"
-
-
-class StatusMessage(BaseModel):
-    """Represents a data model to record success or error when performing registry interactions.
-
-    Args:
-        status (StatusEnum): Represents the outcome (success or error) of the registry interaction.
-        action: (ActionEnum): Represents the type of registry action that was taken.
-        message: (str): Provides a more detailed status message for the registry interaction.
-    """
-
-    status: StatusEnum
-    action: ActionEnum
-    message: str