Rejigger the finalizers and more

Author: abrookins

redisvl/extensions/llmcache/semantic.py

@@ -25,6 +25,7 @@
 from redisvl.redis.connection import RedisConnectionFactory
 from redisvl.utils.log import get_logger
 from redisvl.utils.utils import (
+    sync_wrapper,
     current_timestamp,
     deprecated_argument,
     serialize,
@@ -133,6 +134,12 @@ def __init__(
             name, prefix, vectorizer.dims, vectorizer.dtype  # type: ignore
         )
         schema = self._modify_schema(schema, filterable_fields)
+
+        if redis_client:
+            self._owns_redis_client = False
+        else:
+            self._owns_redis_client = True
+
         self._index = SearchIndex(
             schema=schema,
             redis_client=redis_client,
@@ -153,8 +160,6 @@ 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,
@@ -317,7 +322,9 @@ async def _avectorize_prompt(self, prompt: Optional[str]) -> List[float]:
     def _check_vector_dims(self, vector: List[float]):
         """Checks the size of the provided vector and raises an error if it
         doesn't match the search index vector dimensions."""
-        schema_vector_dims = self._index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.dims  # type: ignore
+        schema_vector_dims = self._index.schema.fields[
+            CACHE_VECTOR_FIELD_NAME
+        ].attrs.dims  # type: ignore
         validate_vector_dims(len(vector), schema_vector_dims)
 
     def check(
@@ -392,7 +399,8 @@ def check(
         # Search the cache!
         cache_search_results = self._index.query(query)
         redis_keys, cache_hits = self._process_cache_results(
-            cache_search_results, return_fields  # type: ignore
+            cache_search_results,
+            return_fields,  # type: ignore
         )
         # Extend TTL on keys
         for key in redis_keys:
@@ -473,7 +481,8 @@ async def acheck(
         # Search the cache!
         cache_search_results = await aindex.query(query)
         redis_keys, cache_hits = self._process_cache_results(
-            cache_search_results, return_fields  # type: ignore
+            cache_search_results,
+            return_fields,  # type: ignore
         )
         # Extend TTL on keys
         await asyncio.gather(*[self._async_refresh_ttl(key) for key in redis_keys])
@@ -646,7 +655,6 @@ def update(self, key: str, **kwargs) -> None:
         """
         if kwargs:
             for k, v in kwargs.items():
-
                 # Make sure the item is in the index schema
                 if k not in set(self._index.schema.field_names + [METADATA_FIELD_NAME]):
                     raise ValueError(f"{k} is not a valid field within the cache entry")
@@ -689,7 +697,6 @@ async def aupdate(self, key: str, **kwargs) -> None:
 
         if kwargs:
             for k, v in kwargs.items():
-
                 # Make sure the item is in the index schema
                 if k not in set(self._index.schema.field_names + [METADATA_FIELD_NAME]):
                     raise ValueError(f"{k} is not a valid field within the cache entry")
@@ -708,29 +715,18 @@ 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._owns_redis_client is False:
+            return
         if self._index:
             self._index.disconnect()
         if self._aindex:
-            asyncio.run(self._aindex.disconnect())
+            self._aindex.disconnect_sync()
 
     async def adisconnect(self):
+        if not self._owns_redis_client:
+            return
         if self._index:
             self._index.disconnect()
         if self._aindex:

redisvl/index/index.py

@@ -1,9 +1,12 @@
 import asyncio
 import json
+import logging
 from os import replace
+from re import S
 import threading
 import warnings
 from functools import wraps
+
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -18,14 +21,13 @@
 )
 import weakref
 
-from redisvl.utils.utils import deprecated_argument, deprecated_function
+from redisvl.utils.utils import deprecated_argument, deprecated_function, sync_wrapper
 
 if TYPE_CHECKING:
     from redis.commands.search.aggregation import AggregateResult
     from redis.commands.search.document import Document
     from redis.commands.search.result import Result
     from redisvl.query.query import BaseQuery
-    import redis.asyncio
 
 import redis
 import redis.asyncio as aredis
@@ -38,7 +40,6 @@
 from redisvl.redis.connection import (
     RedisConnectionFactory,
     convert_index_info_to_schema,
-    validate_modules,
 )
 from redisvl.redis.utils import convert_bytes
 from redisvl.schema import IndexSchema, StorageType
@@ -279,14 +280,21 @@ def __init__(
 
         self._lib_name: Optional[str] = kwargs.pop("lib_name", None)
 
-       # Store connection parameters
+        # Store connection parameters
         self.__redis_client = redis_client
         self._redis_url = redis_url
         self._connection_kwargs = connection_kwargs or {}
-        self._lock = threading.Lock() 
+        self._lock = threading.Lock()
+
+        self._owns_redis_client = redis_client is None
+        if self._owns_redis_client:
+            weakref.finalize(self, self.disconnect)
 
     def disconnect(self):
         """Disconnect from the Redis database."""
+        if self._owns_redis_client is False:
+            print("Index does not own client, not disconnecting")
+            return
         if self.__redis_client:
             self.__redis_client.close()
         self.__redis_client = None
@@ -343,12 +351,12 @@ def from_existing(
     def client(self) -> Optional[redis.Redis]:
         """The underlying redis-py client object."""
         return self.__redis_client
-    
+
     @property
     def _redis_client(self) -> Optional[redis.Redis]:
         """
         Get a Redis client instance.
-        
+
         Lazily creates a Redis client instance if it doesn't exist.
         """
         if self.__redis_client is None:
@@ -359,7 +367,6 @@ def _redis_client(self) -> Optional[redis.Redis]:
                         **self._connection_kwargs,
                     )
         return self.__redis_client
- 
 
     @deprecated_function("connect", "Pass connection parameters in __init__.")
     def connect(self, redis_url: Optional[str] = None, **kwargs):
@@ -371,8 +378,7 @@ def connect(self, redis_url: Optional[str] = None, **kwargs):
 
         Args:
             redis_url (Optional[str], optional): The URL of the Redis server to
-                connect to. If not provided, the method defaults to using the
-                `REDIS_URL` environment variable.
+                connect to.
 
         Raises:
             redis.exceptions.ConnectionError: If the connection to the Redis
@@ -842,9 +848,9 @@ def __init__(
             schema (IndexSchema): Index schema object.
             redis_url (Optional[str], optional): The URL of the Redis server to
                 connect to.
-            redis_client (Optional[aredis.Redis], optional): An
+            redis_client (Optional[aredis.Redis]): An
                 instantiated redis client.
-            connection_kwargs (Dict[str, Any], optional): Redis client connection
+            connection_kwargs (Optional[Dict[str, Any]]): Redis client connection
                 args.
         """
         if "redis_kwargs" in kwargs:
@@ -864,8 +870,9 @@ def __init__(
         self._connection_kwargs = connection_kwargs or {}
         self._lock = asyncio.Lock()
 
-        # Close connections when the object is garbage collected
-        weakref.finalize(self, self._finalize_disconnect)
+        self._owns_redis_client = redis_client is None
+        if self._owns_redis_client:
+            weakref.finalize(self, sync_wrapper(self.disconnect))
 
     @classmethod
     async def from_existing(
@@ -934,7 +941,7 @@ async def connect(self, redis_url: Optional[str] = None, **kwargs):
         await self.set_client(client)
 
     @deprecated_function("set_client", "Pass connection parameters in __init__.")
-    async def set_client(self, redis_client: aredis.Redis):
+    async def set_client(self, redis_client: Union[aredis.Redis, redis.Redis]):
         """
         [DEPRECATED] Manually set the Redis client to use with the search index.
         This method is deprecated; please provide connection parameters in __init__.
@@ -956,16 +963,17 @@ async def _get_client(self) -> aredis.Redis:
                         kwargs["url"] = self._redis_url
                     self._redis_client = (
                         await RedisConnectionFactory._get_aredis_connection(
-                            required_modules=self.required_modules,
-                            **kwargs
+                            required_modules=self.required_modules, **kwargs
                         )
                     )
             await RedisConnectionFactory.validate_async_redis(
                 self._redis_client, self._lib_name
             )
         return self._redis_client
 
-    async def _validate_client(self, redis_client: aredis.Redis) -> aredis.Redis:
+    async def _validate_client(
+        self, redis_client: Union[aredis.Redis, redis.Redis]
+    ) -> aredis.Redis:
         if isinstance(redis_client, redis.Redis):
             warnings.warn(
                 "Converting sync Redis client to async client is deprecated "
@@ -1340,36 +1348,21 @@ 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._owns_redis_client is False:
+            return
         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:
+        if self._redis_client is None or self._owns_redis_client is False:
             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
+        sync_wrapper(self.disconnect)()
 
     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())

redisvl/utils/utils.py

@@ -1,19 +1,23 @@
+import asyncio
 import inspect
 import json
+import logging
 import warnings
 from contextlib import contextmanager
 from enum import Enum
 from functools import wraps
 from time import time
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Coroutine, Dict, Optional
 from warnings import warn
 
 from pydantic import BaseModel
 from ulid import ULID
 
+from redisvl.utils.log import get_logger
+
 
 def create_ulid() -> str:
-    """Generate a unique indentifier to group related Redis documents."""
+    """Generate a unique identifier to group related Redis documents."""
     return str(ULID())
 
 
@@ -159,3 +163,33 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return decorator
+
+
+def sync_wrapper(fn: Callable[[], Coroutine[Any, Any, Any]]) -> Callable[[], None]:
+    def wrapper():
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        try:
+            if loop is None or not loop.is_running():
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+            task = loop.create_task(fn())
+            loop.run_until_complete(task)
+        except RuntimeError:
+            # This could happen if an object stored an event loop and now
+            # that event loop is closed. There's nothing we can do other than
+            # advise the user to use explicit cleanup methods.
+            #
+            # Uses logging module instead of get_logger() to avoid I/O errors
+            # if the wrapped function is called as a finalizer.
+            logging.info(
+                f"Could not run the async function {fn.__name__} because the event loop is closed. "
+                "This usually means the object was not properly cleaned up. Please use explicit "
+                "cleanup methods (e.g., disconnect(), close()) or use the object as an async "
+                "context manager.",
+            )
+            return
+
+    return wrapper