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

aiq/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/aiqtoolkit_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/aiqtoolkit/).

aiq/plugins/redis/memory.py

@@ -0,0 +1,57 @@
+# 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 redis.asyncio as redis
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_memory
+from aiq.data_models.component_ref import EmbedderRef
+from aiq.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="aiq", 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):
+
+    from aiq.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

aiq/plugins/redis/redis_editor.py

@@ -0,0 +1,233 @@
+# 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 aiq.memory.interfaces import MemoryEditor
+from aiq.memory.models import MemoryItem
+
+logger = logging.getLogger(__name__)
+
+INDEX_NAME = "memory_idx"
+
+
+class RedisEditor(MemoryEditor):
+    """
+    Wrapper class that implements AIQ Toolkit 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(f"Attempting to add {len(items)} items to Redis")
+
+        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(f"Generated memory key: {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(f"Prepared memory data for key {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(f"Generated embedding vector of length: {len(search_vector)}")
+                memory_data["embedding"] = search_vector
+
+            try:
+                # Store as JSON in Redis
+                logger.debug(f"Attempting to store memory data in Redis for key: {memory_key}")
+                await self._client.json().set(memory_key, "$", memory_data)
+                logger.debug(f"Successfully stored memory data for key: {memory_key}")
+
+                # Verify the data was stored
+                stored_data = await self._client.json().get(memory_key)
+                logger.debug(f"Verified data storage for key {memory_key}: {bool(stored_data)}")
+
+            except redis_exceptions.ResponseError as e:
+                logger.error(f"Failed to store memory item: {str(e)}")
+                raise
+            except redis_exceptions.ConnectionError as e:
+                logger.error(f"Redis connection error while storing memory item: {str(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(f"Search called with query: {query}, top_k: {top_k}, kwargs: {kwargs}")
+
+        user_id = kwargs.get("user_id", "redis")  # TODO: remove this fallback username
+        logger.debug(f"Using user_id: {user_id}")
+
+        # Perform vector search using Redis search
+        logger.debug("Using embedder for vector search")
+        try:
+            logger.debug(f"Generating embedding for query: '{query}'")
+            query_vector = await self._embedder.aembed_query(query)
+            logger.debug(f"Generated embedding vector of length: {len(query_vector)}")
+        except Exception as e:
+            logger.error(f"Failed to generate embedding: {str(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(f"Created search query: {search_query}")
+        logger.debug(f"Query string: {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(f"Converted vector to bytes of length: {len(query_vector_bytes)}")
+        except Exception as e:
+            logger.error(f"Failed to convert vector to bytes: {str(e)}")
+            raise
+
+        try:
+            # Execute search with vector parameters
+            logger.debug("Executing Redis search with vector parameters")
+            logger.debug(f"Search query parameters: vec length={len(query_vector_bytes)}")
+
+            # Log the actual query being executed
+            logger.debug(f"Full search query: {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(f"Total documents in index: {total_docs.get('num_docs', 0)}")
+            except Exception as e:
+                logger.error(f"Failed to get index info: {str(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(f"Search returned {len(results.docs)} results")
+            logger.debug(f"Total results found: {results.total}")
+
+            # Convert results to MemoryItems
+            memories = []
+            for i, doc in enumerate(results.docs):
+                try:
+                    logger.debug(f"Processing result {i+1}/{len(results.docs)}")
+                    # Get the document data from the correct attribute
+                    memory_data = {
+                        "conversation": getattr(doc, 'conversation', []),
+                        "user_id": getattr(doc, 'user_id', user_id),
+                        "tags": getattr(doc, 'tags', []),
+                        "metadata": getattr(doc, 'metadata', {}),
+                        "memory": getattr(doc, 'memory', "")
+                    }
+                    logger.debug(f"Similarity score: {getattr(doc, 'score', 0)}")
+                    logger.debug(f"Extracted data for result {i+1}: {memory_data}")
+                    memory_item = self._create_memory_item(memory_data, user_id)
+                    memories.append(memory_item)
+                    logger.debug(f"Successfully created MemoryItem for result {i+1}")
+                except Exception as e:
+                    logger.error(f"Failed to process result {i+1}: {str(e)}")
+                    raise
+
+            logger.debug(f"Successfully processed all {len(memories)} results")
+            return memories
+        except redis_exceptions.ResponseError as e:
+            logger.error(f"Search failed with ResponseError: {str(e)}")
+            raise
+        except redis_exceptions.ConnectionError as e:
+            logger.error(f"Search failed with ConnectionError: {str(e)}")
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error during search: {str(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(f"Failed to remove items: {str(e)}")
+            raise
+        except redis_exceptions.ConnectionError as e:
+            logger.error(f"Redis connection error while removing items: {str(e)}")
+            raise

aiq/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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+
+from . import memory

aiq/plugins/redis/schema.py

@@ -0,0 +1,135 @@
+# 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(f"Creating schema with embedding dimension: {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(f"Created embedding field with dimension {embedding_dim}")
+
+    schema = (
+        TextField("$.user_id", as_name="user_id"),
+        TagField("$.tags[*]", as_name="tags"),
+        TextField("$.memory", as_name="memory"),
+        # TextField("$.conversations[*]", as_name="conversations"), # TODO: figure out if/how this should be done
+        embedding_field)
+
+    # Log the schema details
+    logger.info("Schema fields:")
+    for field in schema:
+        logger.info(f"  - {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(f"Checking if index '{INDEX_NAME}' exists...")
+        info = await client.ft(INDEX_NAME).info()
+        logger.info(f"Redis search index '{INDEX_NAME}' exists.")
+
+        # Verify the schema
+        schema = info.get('attributes', [])
+
+        return
+    except redis_exceptions.ResponseError as e:
+        error_msg = str(e)
+        if "no such index" not in error_msg.lower() and "Index needs recreation" not in error_msg:
+            logger.error(f"Unexpected Redis error: {error_msg}")
+            raise
+
+        # Index doesn't exist or needs recreation
+        logger.info(f"Creating Redis search index '{INDEX_NAME}' with prefix '{key_prefix}'")
+
+        # Drop any existing index
+        try:
+            logger.info(f"Attempting to drop existing index '{INDEX_NAME}' if it exists")
+            await client.ft(INDEX_NAME).dropindex()
+            logger.info(f"Successfully dropped existing index '{INDEX_NAME}'")
+        except redis_exceptions.ResponseError as e:
+            if "no such index" not in str(e).lower():
+                logger.warning(f"Error while dropping index: {str(e)}")
+
+        # Create new schema and index
+        schema = create_schema(embedding_dim or DEFAULT_DIM)
+        logger.info(f"Created schema with embedding dimension: {embedding_dim or DEFAULT_DIM}")
+
+        try:
+            # Create the index
+            logger.info(f"Creating new index '{INDEX_NAME}' with schema")
+            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(f"Successfully created Redis search index '{INDEX_NAME}'")
+            logger.debug(f"Redis search index info: {info}")
+
+            # Verify the schema
+            schema = info.get('attributes', [])
+            logger.debug(f"New index schema: {schema}")
+
+        except redis_exceptions.ResponseError as e:
+            logger.error(f"Failed to create index: {str(e)}")
+            raise
+        except redis_exceptions.ConnectionError as e:
+            logger.error(f"Redis connection error while creating index: {str(e)}")
+            raise

nvidia_nat_redis-1.2.0rc5.dist-info/METADATA

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: nvidia-nat-redis
+Version: 1.2.0rc5
+Summary: Subpackage for Redis memory integration in AIQtoolkit
+Keywords: ai,agents,memory
+Classifier: Programming Language :: Python
+Requires-Python: <3.13,>=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: nvidia-nat~=1.2
+Requires-Dist: redis~=4.3.4
+
+<!--
+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/aiqtoolkit_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/aiqtoolkit/).

nvidia_nat_redis-1.2.0rc5.dist-info/RECORD

@@ -0,0 +1,11 @@
+aiq/meta/pypi.md,sha256=tfhFKDhfJtxBVT6J5pnvrIdaBN_Uy5Dy6in6CCZpYpU,1101
+aiq/plugins/redis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aiq/plugins/redis/memory.py,sha256=rz-cbcSOe1knd795783leY4vm0n1C4UtUJtZOHmINcU,2569
+aiq/plugins/redis/redis_editor.py,sha256=2hb5BGSTjWId5sb6dNAA3MLzO6WSVAU-Px-iv7Fzy2Q,10088
+aiq/plugins/redis/register.py,sha256=_ffKNKnMfkB2HzX4Nk_9EW0pwebg3GuzAE-iB-CoC3E,839
+aiq/plugins/redis/schema.py,sha256=lcazZKzWEoJEJmGEt6rk135zJ6kGqw6v2b1B5l1xg7o,5494
+nvidia_nat_redis-1.2.0rc5.dist-info/METADATA,sha256=PrUfA8u5_hdV53lCkhonCJtaZILr-LCmAsxw7ywx5i4,1427
+nvidia_nat_redis-1.2.0rc5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_redis-1.2.0rc5.dist-info/entry_points.txt,sha256=JyY7zDx7SZTvETPEEZuTgDVhe3EN23s1Glyp8RHIoWY,56
+nvidia_nat_redis-1.2.0rc5.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+nvidia_nat_redis-1.2.0rc5.dist-info/RECORD,,

nvidia_nat_redis-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_redis-1.2.0rc5.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[aiq.components]
+aiq_redis = aiq.plugins.redis.register

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

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