fix: Pipeline, async fixes

Author: abrookins

langgraph/checkpoint/redis/aio.py

@@ -10,7 +10,6 @@
 from typing import Any, List, Optional, Sequence, Tuple, Type, cast
 
 from langchain_core.runnables import RunnableConfig
-from redis import WatchError
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import FilterQuery
 from redisvl.query.filter import Num, Tag
@@ -74,6 +73,7 @@ def create_indexes(self) -> None:
 
     async def __aenter__(self) -> AsyncRedisSaver:
         """Async context manager enter."""
+        await self.asetup()
         return self
 
     async def __aexit__(
@@ -83,15 +83,15 @@ async def __aexit__(
         exc_tb: Optional[TracebackType],
     ) -> None:
         """Async context manager exit."""
-        # Close client connections
-        if hasattr(self, "checkpoint_index") and hasattr(
-            self.checkpoint_index, "client"
-        ):
-            await self.checkpoint_index.client.aclose()
-        if hasattr(self, "channel_index") and hasattr(self.channel_index, "client"):
-            await self.channel_index.client.aclose()
-        if hasattr(self, "writes_index") and hasattr(self.writes_index, "client"):
-            await self.writes_index.client.aclose()
+        if self._owns_its_client:
+            await self._redis.aclose()  # type: ignore[attr-defined]
+            await self._redis.connection_pool.disconnect()
+
+            # Prevent RedisVL from attempting to close the client
+            # on an event loop in a separate thread.
+            self.checkpoints_index._redis_client = None
+            self.checkpoint_blobs_index._redis_client = None
+            self.checkpoint_writes_index._redis_client = None
 
     async def asetup(self) -> None:
         """Initialize Redis indexes asynchronously."""
@@ -428,18 +428,24 @@ async def aput_writes(
                     task_id,
                     write_obj["idx"],
                 )
-                async def tx(pipe, key=key, write_obj=write_obj, upsert_case=upsert_case):
+
+                async def tx(
+                    pipe, key=key, write_obj=write_obj, upsert_case=upsert_case
+                ):
                     exists = await pipe.exists(key)
                     if upsert_case:
                         if exists:
-                            await pipe.json().set(key, "$.channel", write_obj["channel"])
+                            await pipe.json().set(
+                                key, "$.channel", write_obj["channel"]
+                            )
                             await pipe.json().set(key, "$.type", write_obj["type"])
                             await pipe.json().set(key, "$.blob", write_obj["blob"])
                         else:
                             await pipe.json().set(key, "$", write_obj)
                     else:
                         if not exists:
                             await pipe.json().set(key, "$", write_obj)
+
                 await self._redis.transaction(tx, key)
 
     def put_writes(
@@ -533,18 +539,12 @@ async def from_conn_string(
         redis_client: Optional[AsyncRedis] = None,
         connection_args: Optional[dict[str, Any]] = None,
     ) -> AsyncIterator[AsyncRedisSaver]:
-        saver: Optional[AsyncRedisSaver] = None
-        try:
-            saver = cls(
-                redis_url=redis_url,
-                redis_client=redis_client,
-                connection_args=connection_args,
-            )
+        async with cls(
+            redis_url=redis_url,
+            redis_client=redis_client,
+            connection_args=connection_args,
+        ) as saver:
             yield saver
-        finally:
-            if saver and saver._owns_its_client:  # Ensure saver is not None
-                await saver._redis.aclose()  # type: ignore[attr-defined]
-                await saver._redis.connection_pool.disconnect()
 
     async def aget_channel_values(
         self, thread_id: str, checkpoint_ns: str = "", checkpoint_id: str = ""

langgraph/checkpoint/redis/ashallow.py

@@ -8,7 +8,6 @@
 from typing import Any, AsyncIterator, Dict, List, Optional, Sequence, Tuple, cast
 
 from langchain_core.runnables import RunnableConfig
-from redis import WatchError
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import FilterQuery
 from redisvl.query.filter import Num, Tag
@@ -100,9 +99,22 @@ def __init__(
             redis_client=redis_client,
             connection_args=connection_args,
         )
-        # self.lock = asyncio.Lock()
         self.loop = asyncio.get_running_loop()
 
+    async def __aenter__(self) -> AsyncShallowRedisSaver:
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        if self._owns_its_client:
+            await self._redis.aclose()  # type: ignore[attr-defined]
+            await self._redis.connection_pool.disconnect()
+
+            # Prevent RedisVL from attempting to close the client
+            # on an event loop in a separate thread.
+            self.checkpoints_index._redis_client = None
+            self.checkpoint_blobs_index._redis_client = None
+            self.checkpoint_writes_index._redis_client = None
+
     @classmethod
     @asynccontextmanager
     async def from_conn_string(
@@ -113,18 +125,12 @@ async def from_conn_string(
         connection_args: Optional[dict[str, Any]] = None,
     ) -> AsyncIterator[AsyncShallowRedisSaver]:
         """Create a new AsyncShallowRedisSaver instance."""
-        saver: Optional[AsyncShallowRedisSaver] = None
-        try:
-            saver = cls(
-                redis_url=redis_url,
-                redis_client=redis_client,
-                connection_args=connection_args,
-            )
+        async with cls(
+            redis_url=redis_url,
+            redis_client=redis_client,
+            connection_args=connection_args,
+        ) as saver:
             yield saver
-        finally:
-            if saver and saver._owns_its_client:
-                await saver._redis.aclose()  # type: ignore[attr-defined]
-                await saver._redis.connection_pool.disconnect()
 
     async def asetup(self) -> None:
         """Initialize Redis indexes asynchronously."""
@@ -397,6 +403,7 @@ async def aput_writes(
                     write_obj["idx"],
                 )
                 if upsert_case:
+
                     async def tx(pipe, key=key, write_obj=write_obj):
                         exists = await pipe.exists(key)
                         if exists:

langgraph/store/redis/aio.py

@@ -194,6 +194,7 @@ def create_indexes(self) -> None:
 
     async def __aenter__(self) -> AsyncRedisStore:
         """Async context manager enter."""
+        await self.setup()
         return self
 
     async def __aexit__(
@@ -291,26 +292,11 @@ async def _batch_get_ops(
     ) -> None:
         """Execute GET operations in batch asynchronously."""
         for query, _, namespace, items in self._get_batch_GET_ops_queries(get_ops):
-            # Use RedisVL AsyncSearchIndex search
-            search_query = FilterQuery(
-                filter_expression=query,
-                return_fields=["id"],  # Just need the document id
-                num_results=len(items),
-            )
-            res = await self.store_index.search(search_query)
-
-            # Use pipeline to get the actual JSON documents
-            pipeline = self._redis.pipeline(transaction=False)
-            doc_ids = []
-            for doc in res.docs:
-                # The id is already in the correct format (store:prefix:key)
-                pipeline.json().get(doc.id)
-                doc_ids.append(doc.id)
-
-            json_docs = await pipeline.execute()
-
-            # Convert to dictionary format
-            key_to_row = {doc["key"]: doc for doc in json_docs if doc}
+            res = await self.store_index.search(Query(query))
+            # Parse JSON from each document
+            key_to_row = {
+                json.loads(doc.json)["key"]: json.loads(doc.json) for doc in res.docs
+            }
 
             for idx, key in items:
                 if key in key_to_row:

tests/test_async_store.py

@@ -501,8 +501,11 @@ async def test_async_store_with_memory_persistence(
         "distance_type": "cosine",
     }
 
-    async with AsyncRedisStore.from_conn_string(redis_url, index=index_config) as store:
+    async with AsyncRedisStore.from_conn_string(
+        redis_url, index=index_config
+    ) as store, AsyncRedisSaver.from_conn_string(redis_url) as checkpointer:
         await store.setup()
+        await checkpointer.asetup()
 
         model = ChatAnthropic(model="claude-3-5-sonnet-20240620")  # type: ignore[call-arg]
 
@@ -532,11 +535,6 @@ def call_model(
         builder.add_node("call_model", call_model)  # type:ignore[arg-type]
         builder.add_edge(START, "call_model")
 
-        checkpointer = None
-        async with AsyncRedisSaver.from_conn_string(redis_url) as cp:
-            await cp.asetup()
-            checkpointer = cp
-
         # Compile graph with store and checkpointer
         graph = builder.compile(checkpointer=checkpointer, store=store)