nvidia-nat-s3 1.1.0a20251020__py3-none-any.whl

nat/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 nat.builder.builder import Builder
+from nat.cli.register_workflow import register_object_store
+from nat.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] = "NAT_S3_OBJECT_STORE_ACCESS_KEY"
+    SECRET_KEY_ENV: ClassVar[str] = "NAT_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 .s3_object_store import S3ObjectStore
+
+    async with S3ObjectStore(**config.model_dump(exclude={"type"}, exclude_none=False)) as store:
+        yield store

nat/plugins/s3/register.py

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

nat/plugins/s3/s3_object_store.py

@@ -0,0 +1,166 @@
+# 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 nat.data_models.object_store import KeyAlreadyExistsError
+from nat.data_models.object_store import NoSuchKeyError
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+logger = logging.getLogger(__name__)
+
+
+class S3ObjectStore(ObjectStore):
+    """
+    S3ObjectStore is an ObjectStore implementation that uses S3 as the underlying storage.
+    """
+
+    def __init__(self,
+                 *,
+                 bucket_name: str,
+                 endpoint_url: str | None,
+                 access_key: str | None,
+                 secret_key: str | None,
+                 region: str | None):
+
+        super().__init__()
+
+        self.bucket_name = bucket_name
+        self.session = aioboto3.Session()
+        self._client: BaseClient | None = None
+        self._client_context = None
+
+        self._client_args: dict = {}
+        if access_key and secret_key:
+            self._client_args["aws_access_key_id"] = access_key
+            self._client_args["aws_secret_access_key"] = secret_key
+        if region:
+            self._client_args["region_name"] = region
+        if endpoint_url:
+            self._client_args["endpoint_url"] = endpoint_url
+
+    async def __aenter__(self) -> "S3ObjectStore":
+
+        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) -> None:
+
+        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",
+                ) from e
+            # 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)) from e
+            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)) from e
+            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.1.0a20251020.dist-info/METADATA

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: nvidia-nat-s3
+Version: 1.1.0a20251020
+Summary: Subpackage for S3-compatible integration in NeMo Agent toolkit
+Author: NVIDIA Corporation
+Maintainer: NVIDIA Corporation
+License: Apache-2.0
+Project-URL: documentation, https://docs.nvidia.com/nemo/agent-toolkit/latest/
+Project-URL: source, https://github.com/NVIDIA/NeMo-Agent-Toolkit
+Keywords: ai,agents,memory,data store
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE-3rd-party.txt
+License-File: LICENSE.md
+Requires-Dist: nvidia-nat==v1.1.0a20251020
+Requires-Dist: aioboto3>=11.0.0
+Dynamic: license-file

nvidia_nat_s3-1.1.0a20251020.dist-info/RECORD

@@ -0,0 +1,11 @@
+nat/plugins/s3/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/s3/object_store.py,sha256=FDcO_8j__Nb-W8IY7JeUk5WoCRIqVNi1a1L1aVmYt10,2172
+nat/plugins/s3/register.py,sha256=iBGq69m0rbZ3BTisA4Lt5UCHtW0Pc7m-l8G8dz-_pBc,813
+nat/plugins/s3/s3_object_store.py,sha256=_2Sx1QXosuS2FnHXtG8qo4x9-xTNJVbNxoiJCZQcWgE,6052
+nvidia_nat_s3-1.1.0a20251020.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat_s3-1.1.0a20251020.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat_s3-1.1.0a20251020.dist-info/METADATA,sha256=dYXVETGEaiIRtkFpDfz39iT98W3YSzUlvc_fbUH2gzQ,823
+nvidia_nat_s3-1.1.0a20251020.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_s3-1.1.0a20251020.dist-info/entry_points.txt,sha256=HSc9lsaEu-3DyVezRMR-VZrfhWnDtA9llVaWE2CYZNw,63
+nvidia_nat_s3-1.1.0a20251020.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_s3-1.1.0a20251020.dist-info/RECORD,,

nvidia_nat_s3-1.1.0a20251020.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.1.0a20251020.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[nat.components]
+nat_s3_object_store = nat.plugins.s3.register