nvidia-nat-redis 1.3.0a20250827__py3-none-any.whl → 1.3.0a20250828__py3-none-any.whl

nat/plugins/redis/memory.py

@@ -13,21 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import redis.asyncio as redis
 from pydantic import Field
 
 from nat.builder.builder import Builder
-from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_memory
 from nat.data_models.component_ref import EmbedderRef
 from nat.data_models.memory import MemoryBaseConfig
 
 
 class RedisMemoryClientConfig(MemoryBaseConfig, name="redis_memory"):
-    host: str | None = Field(default="localhost", description="Redis server host")
-    db: str | None = Field(default="0", description="Redis DB")
-    port: str | None = Field(default="6379", description="Redis server port")
-    key_prefix: str | None = Field(default="nat", description="Key prefix to use for redis keys")
+    host: str = Field(default="localhost", description="Redis server host")
+    db: int = Field(default=0, description="Redis DB")
+    port: int = Field(default=6379, description="Redis server port")
+    key_prefix: str = Field(default="nat", description="Key prefix to use for redis keys")
     embedder: EmbedderRef = Field(description=("Instance name of the memory client instance from the workflow "
                                                "configuration object."))
 
@@ -35,6 +33,9 @@ class RedisMemoryClientConfig(MemoryBaseConfig, name="redis_memory"):
 @register_memory(config_type=RedisMemoryClientConfig)
 async def redis_memory_client(config: RedisMemoryClientConfig, builder: Builder):
 
+    import redis.asyncio as redis
+
+    from nat.builder.framework_enum import LLMFrameworkEnum
     from nat.plugins.redis.redis_editor import RedisEditor
 
     from .schema import ensure_index_exists

nat/plugins/redis/object_store.py

@@ -0,0 +1,40 @@
+# 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.
+
+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 RedisObjectStoreClientConfig(ObjectStoreBaseConfig, name="redis"):
+    """
+    Object store that stores objects in a Redis database.
+    """
+
+    host: str = Field(default="localhost", description="The host of the Redis server")
+    db: int = Field(default=0, description="The Redis logical database number")
+    port: int = Field(default=6379, description="The port of the Redis server")
+    bucket_name: str = Field(description="The name of the bucket to use for the object store")
+
+
+@register_object_store(config_type=RedisObjectStoreClientConfig)
+async def redis_object_store_client(config: RedisObjectStoreClientConfig, _builder: Builder):
+
+    from .redis_object_store import RedisObjectStore
+
+    async with RedisObjectStore(**config.model_dump(exclude={"type"})) as store:
+        yield store

nat/plugins/redis/redis_object_store.py

@@ -0,0 +1,126 @@
+# 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 redis.asyncio as redis
+
+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
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+
+class RedisObjectStore(ObjectStore):
+    """
+    Implementation of ObjectStore that stores objects in Redis.
+
+    Each object is stored as a single binary value at key "nat/object_store/{bucket_name}/{object_key}".
+    """
+
+    def __init__(self, *, bucket_name: str, host: str, port: int, db: int):
+
+        super().__init__()
+
+        self._bucket_name = bucket_name
+        self._host = host
+        self._port = port
+        self._db = db
+        self._client: redis.Redis | None = None
+
+    async def __aenter__(self) -> "RedisObjectStore":
+
+        if self._client is not None:
+            raise RuntimeError("Connection already established")
+
+        self._client = redis.Redis(
+            host=self._host,
+            port=self._port,
+            db=self._db,
+            socket_timeout=5.0,
+            socket_connect_timeout=5.0,
+        )
+
+        # Ping to ensure connectivity
+        res = await self._client.ping()
+        if not res:
+            raise RuntimeError("Failed to connect to Redis")
+
+        logger.info("Connected Redis client for %s at %s:%s/%s", self._bucket_name, self._host, self._port, self._db)
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+
+        if not self._client:
+            raise RuntimeError("Connection not established")
+
+        await self._client.close()
+        self._client = None
+
+    def _make_key(self, key: str) -> str:
+        return f"nat/object_store/{self._bucket_name}/{key}"
+
+    @override
+    async def put_object(self, key: str, item: ObjectStoreItem):
+
+        if not self._client:
+            raise RuntimeError("Connection not established")
+
+        full_key = self._make_key(key)
+
+        item_json = item.model_dump_json()
+        # Redis SET with NX ensures we do not overwrite existing keys
+        if not await self._client.set(full_key, item_json, nx=True):
+            raise KeyAlreadyExistsError(key=key,
+                                        additional_message=f"Redis bucket {self._bucket_name} already has key {key}")
+
+    @override
+    async def upsert_object(self, key: str, item: ObjectStoreItem):
+
+        if not self._client:
+            raise RuntimeError("Connection not established")
+
+        full_key = self._make_key(key)
+        item_json = item.model_dump_json()
+        await self._client.set(full_key, item_json)
+
+    @override
+    async def get_object(self, key: str) -> ObjectStoreItem:
+
+        if not self._client:
+            raise RuntimeError("Connection not established")
+
+        full_key = self._make_key(key)
+        data = await self._client.get(full_key)
+        if data is None:
+            raise NoSuchKeyError(key=key,
+                                 additional_message=f"Redis bucket {self._bucket_name} does not have key {key}")
+        return ObjectStoreItem.model_validate_json(data)
+
+    @override
+    async def delete_object(self, key: str):
+
+        if not self._client:
+            raise RuntimeError("Connection not established")
+
+        full_key = self._make_key(key)
+        deleted = await self._client.delete(full_key)
+        if deleted == 0:
+            raise NoSuchKeyError(key=key,
+                                 additional_message=f"Redis bucket {self._bucket_name} does not have key {key}")

nat/plugins/redis/register.py

@@ -19,3 +19,4 @@
 # Import any providers which need to be automatically registered here
 
 from . import memory
+from . import object_store

{nvidia_nat_redis-1.3.0a20250827.dist-info → nvidia_nat_redis-1.3.0a20250828.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-redis
-Version: 1.3.0a20250827
+Version: 1.3.0a20250828
 Summary: Subpackage for Redis integration in NeMo Agent toolkit
 Keywords: ai,agents,memory
 Classifier: Programming Language :: Python
 Requires-Python: <3.13,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3.0a20250827
+Requires-Dist: nvidia-nat==v1.3.0a20250828
 Requires-Dist: redis~=4.3.4
 
 <!--

nvidia_nat_redis-1.3.0a20250828.dist-info/RECORD

@@ -0,0 +1,13 @@
+nat/meta/pypi.md,sha256=TpeNbVZJxzvEf0Gh3BGvLHPYsKnXjgM_KQVCayBPXso,1090
+nat/plugins/redis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/redis/memory.py,sha256=wqj_UYqnllc4kTCtFDz3eu-OZnyPWXxNZ3e6G992OGQ,2546
+nat/plugins/redis/object_store.py,sha256=f_GtCTZ3KHfxE4f0lAzKaoAoInAZZki4Dfo7hrKCAHA,1681
+nat/plugins/redis/redis_editor.py,sha256=BskJN5R0h1xzz7w5Ic8IVvd4z-qdBp9ZBIpGYNiSwE8,10080
+nat/plugins/redis/redis_object_store.py,sha256=DX46GEQl4H1Ivf2wrRaSimNcq6EfvRsm-xhyiECudoQ,4302
+nat/plugins/redis/register.py,sha256=dJBKi-7W72ipkmZTOIo1E3ETffmJIlYhQTOlrkiFH3A,834
+nat/plugins/redis/schema.py,sha256=lcazZKzWEoJEJmGEt6rk135zJ6kGqw6v2b1B5l1xg7o,5494
+nvidia_nat_redis-1.3.0a20250828.dist-info/METADATA,sha256=pKeK9u7-KKFMn3oiGknZ8_QLPNK9yJFn_EqklMMR2rw,1435
+nvidia_nat_redis-1.3.0a20250828.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_redis-1.3.0a20250828.dist-info/entry_points.txt,sha256=nyS8t8L9CbRFIMlE70RQBtJXrflBP4Ltl5zAkIl44So,56
+nvidia_nat_redis-1.3.0a20250828.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_redis-1.3.0a20250828.dist-info/RECORD,,

nvidia_nat_redis-1.3.0a20250827.dist-info/RECORD

@@ -1,11 +0,0 @@
-nat/meta/pypi.md,sha256=TpeNbVZJxzvEf0Gh3BGvLHPYsKnXjgM_KQVCayBPXso,1090
-nat/plugins/redis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nat/plugins/redis/memory.py,sha256=2AbG29r6tul4RVvGnSVrQnT96Hvm725i1X23zvDGWTg,2569
-nat/plugins/redis/redis_editor.py,sha256=BskJN5R0h1xzz7w5Ic8IVvd4z-qdBp9ZBIpGYNiSwE8,10080
-nat/plugins/redis/register.py,sha256=n5QJ6VHpLSV4j9WJEc7IgvWHbSyfy8KYeE9XKz4NZxo,807
-nat/plugins/redis/schema.py,sha256=lcazZKzWEoJEJmGEt6rk135zJ6kGqw6v2b1B5l1xg7o,5494
-nvidia_nat_redis-1.3.0a20250827.dist-info/METADATA,sha256=wVW_vpHo4X3rP78jmu2o_Wr_4ChjBxnsOLW_mO3Q7RQ,1435
-nvidia_nat_redis-1.3.0a20250827.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_redis-1.3.0a20250827.dist-info/entry_points.txt,sha256=nyS8t8L9CbRFIMlE70RQBtJXrflBP4Ltl5zAkIl44So,56
-nvidia_nat_redis-1.3.0a20250827.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_redis-1.3.0a20250827.dist-info/RECORD,,