add finalizers, explicit disconnect methods, ctx mgrs

Author: abrookins

redisvl/extensions/llmcache/semantic.py

@@ -1,8 +1,8 @@
 import asyncio
 from typing import Any, Dict, List, Optional
+import weakref
 
 from redis import Redis
-from redis.asyncio import Redis as AsyncRedis
 from redisvl.extensions.constants import (
     CACHE_VECTOR_FIELD_NAME,
     ENTRY_ID_FIELD_NAME,
@@ -23,6 +23,7 @@
 from redisvl.query import RangeQuery
 from redisvl.query.filter import FilterExpression
 from redisvl.redis.connection import RedisConnectionFactory
+from redisvl.utils.log import get_logger
 from redisvl.utils.utils import (
     current_timestamp,
     deprecated_argument,
@@ -32,6 +33,9 @@
 from redisvl.utils.vectorize import BaseVectorizer, HFTextVectorizer
 
 
+logger = get_logger("[RedisVL]")
+
+
 class SemanticCache(BaseLLMCache):
     """Semantic Cache for Large Language Models."""
 
@@ -149,6 +153,8 @@ def __init__(
 
         # Create the search index in Redis
         self._index.create(overwrite=overwrite, drop=False)
+        
+        weakref.finalize(self, self._finalize_async)
 
     def _modify_schema(
         self,
@@ -702,3 +708,37 @@ async def aupdate(self, key: str, **kwargs) -> None:
             await aindex.load(data=[kwargs], keys=[key])
 
         await self._async_refresh_ttl(key)
+        
+    def _finalize_async(self):
+        if self._index:
+            self._index.disconnect()
+        if self._aindex:
+            try:
+                loop = None
+                try:
+                    loop = asyncio.get_running_loop()
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                loop.run_until_complete(self._aindex.disconnect())
+            except Exception as e:
+                logger.info(f"Error disconnecting from index: {e}")
+
+    def disconnect(self):
+        if self._index:
+            self._index.disconnect()
+        if self._aindex:
+            asyncio.run(self._aindex.disconnect())
+
+    async def adisconnect(self):
+        if self._index:
+            self._index.disconnect()
+        if self._aindex:
+            await self._aindex.disconnect()
+            self._aindex = None
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.adisconnect()

redisvl/index/index.py

@@ -16,6 +16,7 @@
     Optional,
     Union,
 )
+import weakref
 
 from redisvl.utils.utils import deprecated_argument, deprecated_function
 
@@ -784,6 +785,12 @@ def info(self, name: Optional[str] = None) -> Dict[str, Any]:
         index_name = name or self.schema.index.name
         return self._info(index_name, self._redis_client)  # type: ignore
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.disconnect()
+
 
 class AsyncSearchIndex(BaseSearchIndex):
     """A search index class for interacting with Redis as a vector database in
@@ -857,11 +864,8 @@ def __init__(
         self._connection_kwargs = connection_kwargs or {}
         self._lock = asyncio.Lock()
 
-    async def disconnect(self):
-        """Asynchronously disconnect and cleanup the underlying async redis connection."""
-        if self._redis_client is not None:
-            await self._redis_client.aclose()  # type: ignore
-        self._redis_client = None
+        # Close connections when the object is garbage collected
+        weakref.finalize(self, self._finalize_disconnect)
 
     @classmethod
     async def from_existing(
@@ -1336,9 +1340,36 @@ async def _info(name: str, redis_client: aredis.Redis) -> Dict[str, Any]:
             raise RedisSearchError(
                 f"Error while fetching {name} index info: {str(e)}"
             ) from e
+            
+    async def disconnect(self):
+        """Asynchronously disconnect and cleanup the underlying async redis connection."""
+        if self._redis_client is not None:
+            await self._redis_client.aclose()  # type: ignore
+        self._redis_client = None
+
+    def disconnect_sync(self):
+        """Synchronously disconnect and cleanup the underlying async redis connection."""
+        if self._redis_client is None:
+            return
+        loop = asyncio.get_running_loop()
+        if loop is None or not loop.is_running():
+            asyncio.run(self._redis_client.aclose())  # type: ignore
+        else:
+            loop.create_task(self.disconnect())
+        self._redis_client = None
 
     async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.disconnect()
+
+    def _finalize_disconnect(self):
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        if loop is None or not loop.is_running():
+            asyncio.run(self.disconnect())
+        else:
+            loop.create_task(self.disconnect())

tests/integration/test_async_search_index.py

@@ -313,3 +313,33 @@ async def test_check_index_exists_before_info(async_index):
 
     with pytest.raises(RedisSearchError):
         await async_index.info()
+
+
+@pytest.mark.asyncio
+async def test_search_index_async_context_manager(async_index):
+    async with async_index:
+        await async_index.create(overwrite=True, drop=True)
+        assert async_index._redis_client
+    assert async_index._redis_client is None
+
+
+@pytest.mark.asyncio
+async def test_search_index_context_manager_with_exception(async_index):
+    async with async_index:
+        await async_index.create(overwrite=True, drop=True)
+        raise ValueError("test")
+    assert async_index._redis_client is None
+    
+    
+@pytest.mark.asyncio
+async def test_search_index_disconnect(async_index):
+    await async_index.create(overwrite=True, drop=True)
+    await async_index.disconnect()
+    assert async_index._redis_client is None
+    
+
+@pytest.mark.asyncio
+async def test_search_engine_disconnect_sync(async_index):
+    await async_index.create(overwrite=True, drop=True)
+    async_index.disconnect_sync()
+    assert async_index._redis_client is None

tests/integration/test_llmcache.py

@@ -941,3 +941,36 @@ def test_deprecated_dtype_argument(redis_url):
             redis_url=redis_url,
             overwrite=True,
         )
+
+
+
+@pytest.mark.asyncio
+async def test_cache_async_context_manager(redis_url):
+    async with SemanticCache(name="test_cache", redis_url=redis_url) as cache:
+        await cache.astore("test prompt", "test response")
+        assert cache._aindex
+    assert cache._aindex is None
+
+
+@pytest.mark.asyncio
+async def test_cache_async_context_manager_with_exception(redis_url):
+    async with SemanticCache(name="test_cache", redis_url=redis_url) as cache:
+        await cache.astore("test prompt", "test response")
+        raise ValueError("test")
+    assert cache._aindex is None
+
+
+@pytest.mark.asyncio
+async def test_cache_async_disconnect(redis_url):
+    cache = SemanticCache(name="test_cache", redis_url=redis_url)
+    await cache.astore("test prompt", "test response")
+    await cache.adisconnect()
+    assert cache._aindex is None
+
+
+def test_cache_disconnect(redis_url):
+    cache = SemanticCache(name="test_cache", redis_url=redis_url)
+    cache.store("test prompt", "test response")
+    cache.disconnect()
+    # We keep this index object around because it isn't lazily created
+    assert cache._index.client is None

tests/integration/test_search_index.py

@@ -287,3 +287,24 @@ def test_check_index_exists_before_info(index):
 def test_index_needs_valid_schema():
     with pytest.raises(ValueError, match=r"Must provide a valid IndexSchema object"):
         SearchIndex(schema="Not A Valid Schema")  # type: ignore
+
+
+def test_search_index_context_manager(index):
+    with index:
+        index.create(overwrite=True, drop=True)
+        assert index.client
+    assert index.client is None
+    
+    
+def test_search_index_context_manager_with_exception(index):
+    with pytest.raises(ValueError):
+        with index:
+            index.create(overwrite=True, drop=True)
+            raise ValueError("test")
+    assert index.client is None
+
+
+def test_search_index_disconnect(index):
+    index.create(overwrite=True, drop=True)
+    index.disconnect()
+    assert index.client is None