From 62270bd1456b15a976e72e1107128c95fb903ec0 Mon Sep 17 00:00:00 2001
From: Paul Breuil <tykapl.breuil@gmail.com>
Date: Wed, 8 Jan 2025 15:04:23 +0100
Subject: [PATCH] feat: add ttl support to vector store

---
 libs/redis/langchain_redis/vectorstores.py       | 10 ++++++++--
 libs/redis/tests/unit_tests/test_vectorstores.py |  5 ++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/libs/redis/langchain_redis/vectorstores.py b/libs/redis/langchain_redis/vectorstores.py
index 59ef8e8..b936954 100644
--- a/libs/redis/langchain_redis/vectorstores.py
+++ b/libs/redis/langchain_redis/vectorstores.py
@@ -159,6 +159,8 @@ class RedisVectorStore(VectorStore):
             URL of the Redis instance to connect to.
         redis_client: Optional[Redis]
             Pre-existing Redis connection.
+        ttl: Optional[int]
+            Time-to-live for the Redis keys.
 
     Instantiate:
         .. code-block:: python
@@ -265,10 +267,12 @@ def __init__(
         self,
         embeddings: Embeddings,
         config: Optional[RedisConfig] = None,
+        ttl: Optional[int] = None,
         **kwargs: Any,
     ):
         self.config = config or RedisConfig(**kwargs)
         self._embeddings = embeddings
+        self.ttl = ttl
 
         if self.config.embedding_dimensions is None:
             self.config.embedding_dimensions = len(
@@ -441,10 +445,12 @@ def add_texts(
 
         result = (
             self._index.load(
-                datas, keys=[f"{self.config.key_prefix}:{key}" for key in keys]
+                datas,
+                keys=[f"{self.config.key_prefix}:{key}" for key in keys],
+                ttl=self.ttl,
             )
             if keys
-            else self._index.load(datas)
+            else self._index.load(datas, ttl=self.ttl)
         )
 
         return list(result) if result is not None else []
diff --git a/libs/redis/tests/unit_tests/test_vectorstores.py b/libs/redis/tests/unit_tests/test_vectorstores.py
index 69a24f0..935c79f 100644
--- a/libs/redis/tests/unit_tests/test_vectorstores.py
+++ b/libs/redis/tests/unit_tests/test_vectorstores.py
@@ -70,7 +70,10 @@ def create(self, overwrite: bool = False) -> None:
         pass
 
     def load(
-        self, documents: List[Dict[str, Any]], keys: Optional[List[str]] = None
+        self,
+        documents: List[Dict[str, Any]],
+        keys: Optional[List[str]] = None,
+        ttl: Optional[int] = None,
     ) -> List[str]:
         for i, doc in enumerate(documents):
             key = keys[i] if keys else f"key_{i}"