nvidia-nat-redis 1.3.0a20251112__py3-none-any.whl

nat/meta/pypi.md

@@ -0,0 +1,23 @@
+<!--
+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.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA NeMo Agent Toolkit Subpackage
+This is a subpackage for Redis memory integration in NeMo Agent toolkit.
+
+For more information about NeMo Agent toolkit, please visit the [NeMo Agent toolkit package](https://pypi.org/project/nvidia-nat/).

nat/plugins/redis/memory.py

@@ -0,0 +1,58 @@
+# 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_memory
+from nat.data_models.component_ref import EmbedderRef
+from nat.data_models.memory import MemoryBaseConfig
+
+
+class RedisMemoryClientConfig(MemoryBaseConfig, name="redis_memory"):
+    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."))
+
+
+@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
+
+    redis_client = redis.Redis(host=config.host,
+                               port=config.port,
+                               db=config.db,
+                               decode_responses=True,
+                               socket_timeout=5.0,
+                               socket_connect_timeout=5.0)
+
+    embedder = await builder.get_embedder(config.embedder, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    test_embedding = await embedder.aembed_query("test")
+    embedding_dim = len(test_embedding)
+    await ensure_index_exists(client=redis_client, key_prefix=config.key_prefix, embedding_dim=embedding_dim)
+
+    memory_editor = RedisEditor(redis_client=redis_client, key_prefix=config.key_prefix, embedder=embedder)
+
+    yield memory_editor

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_editor.py

@@ -0,0 +1,228 @@
+# 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 secrets
+
+import numpy as np
+import redis.asyncio as redis
+import redis.exceptions as redis_exceptions
+from langchain_core.embeddings import Embeddings
+from redis.commands.search.query import Query
+
+from nat.memory.interfaces import MemoryEditor
+from nat.memory.models import MemoryItem
+
+logger = logging.getLogger(__name__)
+
+INDEX_NAME = "memory_idx"
+
+
+class RedisEditor(MemoryEditor):
+    """
+    Wrapper class that implements NAT interfaces for Redis memory storage.
+    """
+
+    def __init__(self, redis_client: redis.Redis, key_prefix: str, embedder: Embeddings):
+        """
+        Initialize Redis client for memory storage.
+
+        Args:
+            redis_client: (redis.Redis) Redis client
+            key_prefix: (str) Redis key prefix
+            embedder: (Embeddings) Embedder for semantic search functionality
+        """
+
+        self._client: redis.Redis = redis_client
+        self._key_prefix: str = key_prefix
+        self._embedder: Embeddings = embedder
+
+    async def add_items(self, items: list[MemoryItem]) -> None:
+        """
+        Insert Multiple MemoryItems into Redis.
+        Each MemoryItem is stored with its metadata and tags.
+        """
+        logger.debug("Attempting to add %d items to Redis", len(items))
+
+        for memory_item in items:
+            item_meta = memory_item.metadata
+            conversation = memory_item.conversation
+            user_id = memory_item.user_id
+            tags = memory_item.tags
+            memory_id = secrets.token_hex(4)  # e.g. 02ba3fe9
+
+            # Create a unique key for this memory item
+            memory_key = f"{self._key_prefix}:memory:{memory_id}"
+            logger.debug("Generated memory key: %s", memory_key)
+
+            # Prepare memory data
+            memory_data = {
+                "conversation": conversation,
+                "user_id": user_id,
+                "tags": tags,
+                "metadata": item_meta,
+                "memory": memory_item.memory or ""
+            }
+            logger.debug("Prepared memory data for key %s", memory_key)
+
+            # If we have memory, compute and store the embedding
+            if memory_item.memory:
+                logger.debug("Computing embedding for memory text")
+                search_vector = await self._embedder.aembed_query(memory_item.memory)
+                logger.debug("Generated embedding vector of length: %d", len(search_vector))
+                memory_data["embedding"] = search_vector
+
+            try:
+                # Store as JSON in Redis
+                logger.debug("Attempting to store memory data in Redis for key: %s", memory_key)
+                await self._client.json().set(memory_key, "$", memory_data)
+                logger.debug("Successfully stored memory data for key: %s", memory_key)
+
+                # Verify the data was stored
+                stored_data = await self._client.json().get(memory_key)
+                logger.debug("Verified data storage for key %s: %s", memory_key, bool(stored_data))
+
+            except redis_exceptions.ResponseError as e:
+                logger.error("Failed to store memory item: %s", e)
+                raise
+            except redis_exceptions.ConnectionError as e:
+                logger.error("Redis connection error while storing memory item: %s", e)
+                raise
+
+    async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:
+        """
+        Retrieve items relevant to the given query.
+
+        Args:
+            query (str): The query string to match.
+            top_k (int): Maximum number of items to return.
+            kwargs (dict): Keyword arguments to pass to the search method.
+
+        Returns:
+            list[MemoryItem]: The most relevant MemoryItems for the given query.
+        """
+        logger.debug("Search called with query: %s, top_k: %d, kwargs: %s", query, top_k, kwargs)
+
+        user_id = kwargs.get("user_id", "redis")  # TODO: remove this fallback username
+        logger.debug("Using user_id: %s", user_id)
+
+        # Perform vector search using Redis search
+        logger.debug("Using embedder for vector search")
+        try:
+            logger.debug("Generating embedding for query: '%s'", query)
+            query_vector = await self._embedder.aembed_query(query)
+            logger.debug("Generated embedding vector of length: %d", len(query_vector))
+        except Exception as e:
+            logger.error("Failed to generate embedding: %s", e)
+            raise
+
+        # Create vector search query
+        search_query = (
+            Query(f"(@user_id:{user_id})=>[KNN {top_k} @embedding $vec AS score]").sort_by("score").return_fields(
+                "conversation", "user_id", "tags", "metadata", "memory", "score").dialect(2))
+        logger.debug("Created search query: %s", search_query)
+        logger.debug("Query string: %s", search_query.query_string())
+
+        # Convert query vector to bytes
+        try:
+            logger.debug("Converting query vector to bytes")
+            query_vector_bytes = np.array(query_vector, dtype=np.float32).tobytes()
+            logger.debug("Converted vector to bytes of length: %d", len(query_vector_bytes))
+        except Exception as e:
+            logger.error("Failed to convert vector to bytes: %s", e)
+            raise
+
+        try:
+            # Execute search with vector parameters
+            logger.debug("Executing Redis search with vector parameters")
+            logger.debug("Search query parameters: vec length=%d", len(query_vector_bytes))
+
+            # Log the actual query being executed
+            logger.debug("Full search query: %s", search_query.query_string())
+
+            # Check if there are any documents in the index
+            try:
+                total_docs = await self._client.ft(INDEX_NAME).info()
+                logger.debug("Total documents in index: %d", total_docs.get('num_docs', 0))
+            except Exception as e:
+                logger.exception("Failed to get index info: %s", e)
+
+            # Execute the search
+            results = await self._client.ft(INDEX_NAME).search(search_query, query_params={"vec": query_vector_bytes})
+
+            # Log detailed results information
+            logger.debug("Search returned %d results", len(results.docs))
+            logger.debug("Total results found: %d", results.total)
+
+            # Convert results to MemoryItems
+            memories = []
+            for i, doc in enumerate(results.docs):
+                try:
+                    logger.debug("Processing result %d/%d", i + 1, len(results.docs))
+                    logger.debug("Similarity score: %d", getattr(doc, 'score', 0))
+
+                    # Get the full document data
+                    full_doc = await self._client.json().get(doc.id)
+                    logger.debug("Extracted data for result %d: %s", i + 1, full_doc)
+                    memory_item = self._create_memory_item(dict(full_doc), user_id)
+                    memories.append(memory_item)
+                    logger.debug("Successfully created MemoryItem for result %d", i + 1)
+                except Exception as e:
+                    logger.error("Failed to process result %d: %s", i + 1, e)
+                    raise
+
+            logger.debug("Successfully processed all %d results", len(memories))
+            return memories
+        except redis_exceptions.ResponseError as e:
+            logger.error("Search failed with ResponseError: %s", e)
+            raise
+        except redis_exceptions.ConnectionError as e:
+            logger.error("Search failed with ConnectionError: %s", e)
+            raise
+        except Exception as e:
+            logger.error("Unexpected error during search: %s", e)
+            raise
+
+    def _create_memory_item(self, memory_data: dict, user_id: str) -> MemoryItem:
+        """Helper method to create a MemoryItem from Redis data."""
+        # Ensure tags is always a list
+        tags = memory_data.get("tags", [])
+        # Not sure why but sometimes the tags are retrieved as a string
+        if isinstance(tags, str):
+            tags = [tags]
+        elif not isinstance(tags, list):
+            tags = []
+
+        return MemoryItem(conversation=memory_data.get("conversation", []),
+                          user_id=user_id,
+                          memory=memory_data.get("memory", ""),
+                          tags=tags,
+                          metadata=memory_data.get("metadata", {}))
+
+    async def remove_items(self, **kwargs):
+        """
+        Remove memory items based on provided criteria.
+        """
+        try:
+            pattern = f"{self._key_prefix}:memory:*"
+            keys = await self._client.keys(pattern)
+            if keys:
+                await self._client.delete(*keys)
+        except redis_exceptions.ResponseError as e:
+            logger.error("Failed to remove items: %s", e)
+            raise
+        except redis_exceptions.ConnectionError as e:
+            logger.error("Redis connection error while removing items: %s", e)
+            raise

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

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

nat/plugins/redis/schema.py

@@ -0,0 +1,136 @@
+# 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
+import redis.exceptions as redis_exceptions
+from redis.commands.search.field import TagField
+from redis.commands.search.field import TextField
+from redis.commands.search.field import VectorField
+from redis.commands.search.indexDefinition import IndexDefinition
+from redis.commands.search.indexDefinition import IndexType
+
+logger = logging.getLogger(__name__)
+
+INDEX_NAME = "memory_idx"
+DEFAULT_DIM = 384  # Default embedding dimension
+
+
+def create_schema(embedding_dim: int = DEFAULT_DIM):
+    """
+    Create the Redis search schema for redis_memory.
+
+    Args:
+        embedding_dim (int): Dimension of the embedding vectors
+
+    Returns:
+        tuple: Schema definition for Redis search
+    """
+    logger.info("Creating schema with embedding dimension: %d", embedding_dim)
+
+    embedding_field = VectorField("$.embedding",
+                                  "HNSW",
+                                  {
+                                      "TYPE": "FLOAT32",
+                                      "DIM": embedding_dim,
+                                      "DISTANCE_METRIC": "L2",
+                                      "INITIAL_CAP": 100,
+                                      "M": 16,
+                                      "EF_CONSTRUCTION": 200,
+                                      "EF_RUNTIME": 10
+                                  },
+                                  as_name="embedding")
+    logger.info("Created embedding field with dimension %d", embedding_dim)
+
+    schema = (
+        # Redis search can't directly index complex objects (e.g. conversation and metadata) in return_fields
+        # They need to be retrieved via json().get() for full object access
+        TextField("$.user_id", as_name="user_id"),
+        TagField("$.tags[*]", as_name="tags"),
+        TextField("$.memory", as_name="memory"),
+        embedding_field)
+
+    # Log the schema details
+    logger.info("Schema fields:")
+    for field in schema:
+        logger.info("  - %s: %s", field.name, type(field).__name__)
+
+    return schema
+
+
+async def ensure_index_exists(client: redis.Redis, key_prefix: str, embedding_dim: int | None) -> None:
+    """
+    Ensure the Redis search index exists, creating it if necessary.
+
+    Args:
+        client (redis.Redis): Redis client instance
+        key_prefix (str): Prefix for keys to be indexed
+        embedding_dim (Optional[int]): Dimension of embedding vectors. If None, uses default.
+    """
+    try:
+        # Check if index exists
+        logger.info("Checking if index '%s' exists...", INDEX_NAME)
+        info = await client.ft(INDEX_NAME).info()
+        logger.info("Redis search index '%s' exists.", INDEX_NAME)
+
+        # Verify the schema
+        schema = info.get('attributes', [])
+
+        return
+    except redis_exceptions.ResponseError as ex:
+        error_msg = str(ex)
+        if "no such index" not in error_msg.lower() and "Index needs recreation" not in error_msg:
+            logger.error("Unexpected Redis error: %s", error_msg)
+            raise
+
+        # Index doesn't exist or needs recreation
+        logger.info("Creating Redis search index '%s' with prefix '%s'", INDEX_NAME, key_prefix)
+
+        # Drop any existing index
+        try:
+            logger.info("Attempting to drop existing index '%s' if it exists", INDEX_NAME)
+            await client.ft(INDEX_NAME).dropindex()
+            logger.info("Successfully dropped existing index '%s'", INDEX_NAME)
+        except redis_exceptions.ResponseError as e:
+            if "no such index" not in str(e).lower():
+                logger.warning("Error while dropping index: %s", str(e))
+
+        # Create new schema and index
+        schema = create_schema(embedding_dim or DEFAULT_DIM)
+        logger.info("Created schema with embedding dimension: %d", embedding_dim or DEFAULT_DIM)
+
+        try:
+            # Create the index
+            logger.info("Creating new index '%s' with schema", INDEX_NAME)
+            await client.ft(INDEX_NAME).create_index(schema,
+                                                     definition=IndexDefinition(prefix=[f"{key_prefix}:"],
+                                                                                index_type=IndexType.JSON))
+
+            # Verify index was created
+            info = await client.ft(INDEX_NAME).info()
+            logger.info("Successfully created Redis search index '%s'", INDEX_NAME)
+            logger.debug("Redis search index info: %s", info)
+
+            # Verify the schema
+            schema = info.get('attributes', [])
+            logger.debug("New index schema: %s", schema)
+
+        except redis_exceptions.ResponseError as e:
+            logger.error("Failed to create index: %s", str(e))
+            raise
+        except redis_exceptions.ConnectionError as e:
+            logger.error("Redis connection error while creating index: %s", str(e))
+            raise