aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

aiq/object_store/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 Toolkit Object Store Module
+
+This package provides foundational classes and interfaces
+for managing object storage in AIQ Toolkit's LLM-based agents.
+"""

aiq/object_store/in_memory_object_store.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_object_store
+from aiq.data_models.object_store import KeyAlreadyExistsError
+from aiq.data_models.object_store import NoSuchKeyError
+from aiq.data_models.object_store import ObjectStoreBaseConfig
+from aiq.utils.type_utils import override
+
+from .interfaces import ObjectStore
+from .models import ObjectStoreItem
+
+
+class InMemoryObjectStoreConfig(ObjectStoreBaseConfig, name="in_memory"):
+    """
+    Object store that stores objects in memory. Objects are not persisted when the process shuts down.
+    """
+    pass
+
+
+class InMemoryObjectStore(ObjectStore):
+    """
+    Implementation of ObjectStore that stores objects in memory. Objects are not persisted when the process shuts down.
+    """
+
+    def __init__(self) -> None:
+        self._store: dict[str, ObjectStoreItem] = {}
+
+    @override
+    async def put_object(self, key: str, item: ObjectStoreItem) -> None:
+        if key in self._store:
+            raise KeyAlreadyExistsError(key)
+
+        self._store[key] = item
+        return
+
+    @override
+    async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
+        self._store[key] = item
+        return
+
+    @override
+    async def get_object(self, key: str) -> ObjectStoreItem:
+
+        if key not in self._store:
+            raise NoSuchKeyError(key)
+
+        return self._store[key]
+
+    @override
+    async def delete_object(self, key: str) -> None:
+        if key not in self._store:
+            raise NoSuchKeyError(key)
+
+        self._store.pop(key)
+        return
+
+
+@register_object_store(config_type=InMemoryObjectStoreConfig)
+async def in_memory_object_store(config: InMemoryObjectStoreConfig, builder: Builder):
+    yield InMemoryObjectStore()

aiq/object_store/interfaces.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 .models import ObjectStoreItem
+
+
+class ObjectStore(ABC):
+    """
+    Abstract interface for an object store.
+
+    Implementations may integrate with various object stores,
+    such as S3, MySQL, etc.
+    """
+
+    @abstractmethod
+    async def put_object(self, key: str, item: ObjectStoreItem) -> None:
+        """
+        Save an ObjectStoreItem in the object store with the given key.
+        If the key already exists, raise an error.
+
+        Args:
+            key (str): The key to save the item under.
+            item (ObjectStoreItem): The item to save.
+
+        Raises:
+            KeyAlreadyExistsError: If the key already exists.
+        """
+        pass
+
+    @abstractmethod
+    async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
+        """
+        Save an ObjectStoreItem in the object store with the given key.
+        If the key already exists, update the item.
+
+        Args:
+            key (str): The key to save the item under.
+            item (ObjectStoreItem): The item to save.
+        """
+        pass
+
+    @abstractmethod
+    async def get_object(self, key: str) -> ObjectStoreItem:
+        """
+        Get an ObjectStoreItem from the object store by key.
+
+        Args:
+            key (str): The key to get the item from.
+
+        Returns:
+            ObjectStoreItem: The item retrieved from the object store.
+
+        Raises:
+            NoSuchKeyError: If the item does not exist.
+        """
+        pass
+
+    @abstractmethod
+    async def delete_object(self, key: str) -> None:
+        """
+        Delete an ObjectStoreItem from the object store by key.
+
+        Args:
+            key (str): The key to delete the item from.
+
+        Raises:
+            NoSuchKeyError: If the item does not exist.
+        """
+        pass

aiq/object_store/models.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 BaseModel
+from pydantic import Field
+
+
+class ObjectStoreItem(BaseModel):
+    """
+    Represents an object store item consisting of bytes and associated metadata.
+
+    Attributes
+    ----------
+    data : bytes
+        The data to store in the object store.
+    content_type : str | None
+        The content type of the data.
+    metadata : dict[str, str] | None
+        Metadata providing context and utility for management operations.
+    """
+
+    data: bytes = Field(description="The data to store in the object store.")
+    content_type: str | None = Field(description="The content type of the data.", default=None)
+    metadata: dict[str, str] | None = Field(description="The metadata of the data.", default=None)

aiq/object_store/register.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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
+# isort:skip_file
+
+from . import in_memory_object_store

aiq/observability/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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/observability/exporter/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.