nvidia-nat-s3 1.2.0rc5__py3-none-any.whl

aiq/plugins/s3/object_store.py

@@ -0,0 +1,49 @@
+# 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 typing import ClassVar
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_object_store
+from aiq.data_models.object_store import ObjectStoreBaseConfig
+
+
+class S3ObjectStoreClientConfig(ObjectStoreBaseConfig, name="s3"):
+    """
+    Object store that stores objects in an S3 bucket.
+    """
+
+    ACCESS_KEY_ENV: ClassVar[str] = "AIQ_S3_OBJECT_STORE_ACCESS_KEY"
+    SECRET_KEY_ENV: ClassVar[str] = "AIQ_S3_OBJECT_STORE_SECRET_KEY"
+
+    bucket_name: str = Field(..., description="The name of the bucket to use for the object store")
+    endpoint_url: str | None = Field(default=None, description="The URL of the S3 server to connect to")
+    access_key: str | None = Field(default=os.environ.get(ACCESS_KEY_ENV),
+                                   description=f"Access key. If omitted, reads from {ACCESS_KEY_ENV}")
+    secret_key: str | None = Field(default=os.environ.get(SECRET_KEY_ENV),
+                                   description=f"Secret key. If omitted, reads from {SECRET_KEY_ENV}")
+    region: str | None = Field(default=None, description="Region to access (or none if unspecified)")
+
+
+@register_object_store(config_type=S3ObjectStoreClientConfig)
+async def s3_object_store_client(config: S3ObjectStoreClientConfig, builder: Builder):
+
+    from aiq.plugins.s3.s3_object_store import S3ObjectStore
+
+    async with S3ObjectStore(config) as store:
+        yield store

aiq/plugins/s3/register.py

@@ -0,0 +1,22 @@
+# 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
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+
+from . import object_store

aiq/plugins/s3/s3_object_store.py

@@ -0,0 +1,168 @@
+# 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.
+
+import logging
+
+import aioboto3
+from botocore.client import BaseClient
+from botocore.exceptions import ClientError
+
+from aiq.data_models.object_store import KeyAlreadyExistsError
+from aiq.data_models.object_store import NoSuchKeyError
+from aiq.object_store.interfaces import ObjectStore
+from aiq.object_store.models import ObjectStoreItem
+from aiq.plugins.s3.object_store import S3ObjectStoreClientConfig
+
+logger = logging.getLogger(__name__)
+
+
+class S3ObjectStore(ObjectStore):
+    """
+    S3ObjectStore is an ObjectStore implementation that uses S3 as the underlying storage.
+    """
+
+    def __init__(self, config: S3ObjectStoreClientConfig):
+
+        super().__init__()
+
+        self.bucket_name = config.bucket_name
+        self.session = aioboto3.Session()
+        self._client: BaseClient | None = None
+        self._client_context = None
+
+        if not config.access_key:
+            raise ValueError("Access key is not set. Please specify it in the environment variable "
+                             "'{S3ObjectStoreClientConfig.ACCESS_KEY_ENV}'.")
+
+        if not config.secret_key:
+            raise ValueError("Secret key is not set. Please specify it in the environment variable "
+                             "'{S3ObjectStoreClientConfig.SECRET_KEY_ENV}'.")
+
+        self._client_args = {
+            "aws_access_key_id": config.access_key,
+            "aws_secret_access_key": config.secret_key,
+            "region_name": config.region,
+            "endpoint_url": config.endpoint_url
+        }
+
+    async def __aenter__(self):
+
+        if self._client_context is not None:
+            raise RuntimeError("Connection already established")
+
+        self._client_context = self.session.client("s3", **self._client_args)
+        if self._client_context is None:
+            raise RuntimeError("Connection unable to be established")
+        self._client = await self._client_context.__aenter__()
+        if self._client is None:
+            raise RuntimeError("Connection unable to be established")
+
+        # Ensure the bucket exists
+        try:
+            await self._client.head_bucket(Bucket=self.bucket_name)
+        except ClientError as e:
+            if e.response['Error']['Code'] == '404':
+                await self._client.create_bucket(Bucket=self.bucket_name)
+                logger.info("Created bucket %s", self.bucket_name)
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+
+        if self._client_context is None:
+            raise RuntimeError("Connection not established")
+
+        await self._client_context.__aexit__(None, None, None)
+        self._client = None
+        self._client_context = None
+
+    async def put_object(self, key: str, item: ObjectStoreItem) -> None:
+
+        if self._client is None:
+            raise RuntimeError("Connection not established")
+
+        put_args = {
+            "Bucket": self.bucket_name,
+            "Key": key,
+            "Body": item.data,
+        }
+        if item.content_type:
+            put_args["ContentType"] = item.content_type
+
+        if item.metadata:
+            put_args["Metadata"] = item.metadata
+
+        try:
+            await self._client.put_object(
+                **put_args,
+                IfNoneMatch='*'  # only succeed if the key does not already exist
+            )
+        except ClientError as e:
+            http_status_code = e.response.get("ResponseMetadata", {}).get("HTTPStatusCode", None)
+            if http_status_code == 412:
+                raise KeyAlreadyExistsError(key=key,
+                                            additional_message=f"S3 object {self.bucket_name}/{key} already exists")
+            else:
+                # Other errors — rethrow or handle accordingly
+                raise
+
+    async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
+
+        if self._client is None:
+            raise RuntimeError("Connection not established")
+
+        put_args = {
+            "Bucket": self.bucket_name,
+            "Key": key,
+            "Body": item.data,
+        }
+        if item.content_type:
+            put_args["ContentType"] = item.content_type
+
+        if item.metadata:
+            put_args["Metadata"] = item.metadata
+
+        await self._client.put_object(**put_args)
+
+    async def get_object(self, key: str) -> ObjectStoreItem:
+        if self._client is None:
+            raise RuntimeError("Connection not established")
+
+        try:
+            response = await self._client.get_object(Bucket=self.bucket_name, Key=key)
+            data = await response["Body"].read()
+            return ObjectStoreItem(data=data, content_type=response['ContentType'], metadata=response['Metadata'])
+        except ClientError as e:
+            if e.response['Error']['Code'] == 'NoSuchKey':
+                raise NoSuchKeyError(key=key, additional_message=str(e))
+            else:
+                raise
+
+    async def delete_object(self, key: str) -> None:
+        if self._client is None:
+            raise RuntimeError("Connection not established")
+
+        try:
+            await self._client.get_object(Bucket=self.bucket_name, Key=key)
+        except ClientError as e:
+            if e.response['Error']['Code'] == 'NoSuchKey':
+                raise NoSuchKeyError(key=key, additional_message=str(e))
+            else:
+                raise
+
+        results = await self._client.delete_object(Bucket=self.bucket_name, Key=key)
+
+        if results.get('DeleteMarker', False):
+            raise NoSuchKeyError(key=key, additional_message="Object was a delete marker")

nvidia_nat_s3-1.2.0rc5.dist-info/METADATA

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: nvidia-nat-s3
+Version: 1.2.0rc5
+Summary: Subpackage for S3-compatible object store integration in AIQ toolkit
+Keywords: ai,agents,memory,data store
+Classifier: Programming Language :: Python
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: nvidia-nat
+Requires-Dist: aioboto3>=11.0.0

nvidia_nat_s3-1.2.0rc5.dist-info/RECORD

@@ -0,0 +1,9 @@
+aiq/plugins/s3/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aiq/plugins/s3/object_store.py,sha256=vVJPxjy2yE55AcUjm_DNyOq7oQGHLiPQ2nohKZRnzog,2134
+aiq/plugins/s3/register.py,sha256=7gqnwyDrYttIlEaa7lo9AASYt-2GrZJE0YT2jpKjepo,845
+aiq/plugins/s3/s3_object_store.py,sha256=l6afjfqRysKKYJKbbxIZ9srSAxUc6XNlT-HEmFI4gUU,6250
+nvidia_nat_s3-1.2.0rc5.dist-info/METADATA,sha256=AxtiVJHO0BAazni5JoU_RY2aq3h9wwJ-L4ZuFSVqDbw,341
+nvidia_nat_s3-1.2.0rc5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_s3-1.2.0rc5.dist-info/entry_points.txt,sha256=7vxhbJ1M2Zz1qFy0ecACW4jidzDMesNC96Ou5p0G7no,63
+nvidia_nat_s3-1.2.0rc5.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+nvidia_nat_s3-1.2.0rc5.dist-info/RECORD,,

nvidia_nat_s3-1.2.0rc5.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nvidia_nat_s3-1.2.0rc5.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[aiq.components]
+aiq_s3_object_store = aiq.plugins.s3.register

nvidia_nat_s3-1.2.0rc5.dist-info/top_level.txt

@@ -0,0 +1 @@
+aiq