Merge pull request #4 from Couchbase-Ecosystem/DA-1012-langgraph-checkpointer

azaddhirajkumar · web-flow · commit eaad9234f162 · 2025-08-14T14:35:17.000+05:30
Create Collection support in Async Class
diff --git a/README.md b/README.md
@@ -68,58 +68,72 @@ with CouchbaseSaver.from_conn_info(
         bucket_name=os.getenv("CB_BUCKET") or "test",
         scope_name=os.getenv("CB_SCOPE") or "langgraph",
     ) as checkpointer:
-        # Create the agent with checkpointing
-        graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
-        
-        # Configure with a unique thread ID
-        config = {"configurable": {"thread_id": "1"}}
-        
-        # Run the agent
-        res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)
-        
-        # Retrieve checkpoints
-        latest_checkpoint = checkpointer.get(config)
-        latest_checkpoint_tuple = checkpointer.get_tuple(config)
-        checkpoint_tuples = list(checkpointer.list(config))
-
-        print(latest_checkpoint)
-        print(latest_checkpoint_tuple)
-        print(checkpoint_tuples)
+    # Create the agent with checkpointing
+    graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
+    
+    # Configure with a unique thread ID
+    config = {"configurable": {"thread_id": "1"}}
+    
+    # Run the agent
+    res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)
+    
+    # Retrieve checkpoints
+    latest_checkpoint = checkpointer.get(config)
+    latest_checkpoint_tuple = checkpointer.get_tuple(config)
+    checkpoint_tuples = list(checkpointer.list(config))
+
+    print(latest_checkpoint)
+    print(latest_checkpoint_tuple)
+    print(checkpoint_tuples)
 ```
 
 ### Asynchronous Usage
 
 ```python
 import os
+from acouchbase.cluster import Cluster as ACluster
+from couchbase.auth import PasswordAuthenticator
+from couchbase.options import ClusterOptions
 from langgraph_checkpointer_couchbase import AsyncCouchbaseSaver
 from langgraph.graph import create_react_agent
 
-async with AsyncCouchbaseSaver.from_conn_info(
-        cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
-        cb_username=os.getenv("CB_USERNAME") or "Administrator",
-        cb_password=os.getenv("CB_PASSWORD") or "password",
-        bucket_name=os.getenv("CB_BUCKET") or "test",
-        scope_name=os.getenv("CB_SCOPE") or "langgraph",
+auth = PasswordAuthenticator(
+    os.getenv("CB_USERNAME") or "Administrator",
+    os.getenv("CB_PASSWORD") or "password",
+)
+options = ClusterOptions(auth)
+cluster = await ACluster.connect(os.getenv("CB_CLUSTER") or "couchbase://localhost", options)
+
+bucket_name = os.getenv("CB_BUCKET") or "test"
+scope_name = os.getenv("CB_SCOPE") or "langgraph"
+
+async with AsyncCouchbaseSaver.from_cluster(
+        cluster=cluster,
+        bucket_name=bucket_name,
+        scope_name=scope_name,
     ) as checkpointer:
-        # Create the agent with checkpointing
-        graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
-        
-        # Configure with a unique thread ID
-        config = {"configurable": {"thread_id": "2"}}
-        
-        # Run the agent asynchronously
-        res = await graph.ainvoke(
-            {"messages": [("human", "what's the weather in nyc")]}, config
-        )
-
-        # Retrieve checkpoints asynchronously
-        latest_checkpoint = await checkpointer.aget(config)
-        latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
-        checkpoint_tuples = [c async for c in checkpointer.alist(config)]
-
-        print(latest_checkpoint)
-        print(latest_checkpoint_tuple)
-        print(checkpoint_tuples)
+    # Create the agent with checkpointing
+    graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
+    
+    # Configure with a unique thread ID
+    config = {"configurable": {"thread_id": "2"}}
+    
+    # Run the agent asynchronously
+    res = await graph.ainvoke(
+        {"messages": [("human", "what's the weather in nyc")]}, config
+    )
+
+    # Retrieve checkpoints asynchronously
+    latest_checkpoint = await checkpointer.aget(config)
+    latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
+    checkpoint_tuples = [c async for c in checkpointer.alist(config)]
+
+    print(latest_checkpoint)
+    print(latest_checkpoint_tuple)
+    print(checkpoint_tuples)
+
+# Close the cluster when done
+await cluster.close()
 ```
 
 ## Configuration Options
diff --git a/langgraph_checkpointer_couchbase/async_cb_saver.py b/langgraph_checkpointer_couchbase/async_cb_saver.py
@@ -1,12 +1,14 @@
 from contextlib import asynccontextmanager
 from datetime import timedelta
 from typing import Any, AsyncIterator, Dict, Optional, Sequence, Tuple
+import logging
 
 from langchain_core.runnables import RunnableConfig
 from acouchbase.cluster import Cluster as ACluster
 from acouchbase.bucket import Bucket as ABucket
 from couchbase.auth import PasswordAuthenticator
 from couchbase.options import ClusterOptions, QueryOptions, UpsertOptions
+from couchbase.exceptions import CollectionAlreadyExistsException
 
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
@@ -18,6 +20,8 @@
 )
 from .utils import _encode_binary, _decode_binary
 
+logger = logging.getLogger(__name__)
+
 class AsyncCouchbaseSaver(BaseCheckpointSaver):
     """A checkpoint saver that stores checkpoints in a Couchbase database."""
 
@@ -35,9 +39,35 @@ def __init__(
         self.cluster = cluster
         self.bucket_name = bucket_name
         self.scope_name = scope_name
+        self.bucket = self.cluster.bucket(bucket_name)
+        self.scope = self.bucket.scope(scope_name)
         self.checkpoints_collection_name = checkpoints_collection_name
         self.checkpoint_writes_collection_name = checkpoint_writes_collection_name
 
+    async def create_collections(self):
+        """Create collections in the Couchbase bucket if they do not exist."""
+
+        collection_manager = self.bucket.collections()
+        try:
+            await collection_manager.create_collection(self.scope_name, self.checkpoints_collection_name)
+        except CollectionAlreadyExistsException as _:
+            pass
+        except Exception as e:
+            logger.exception("Error creating collections")
+            raise e
+        finally:
+            self.checkpoints_collection =  self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
+        
+        try:
+            await collection_manager.create_collection(self.scope_name, self.checkpoint_writes_collection_name)
+        except CollectionAlreadyExistsException as _:
+            pass
+        except Exception as e:
+            logger.exception("Error creating collections")
+            raise e
+        finally:
+            self.checkpoint_writes_collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
+
     @classmethod
     @asynccontextmanager
     async def from_conn_info(
@@ -69,15 +99,25 @@ async def from_conn_info(
             cls.bucket_name = bucket_name
             cls.scope_name = scope_name
 
-            saver = AsyncCouchbaseSaver(cluster, bucket_name, scope_name, checkpoints_collection_name, checkpoint_writes_collection_name)
-            cls.bucket = cluster.bucket(bucket_name)
-            await cls.bucket.on_connect()
+            bucket = cluster.bucket(bucket_name)
+            await bucket.on_connect()
+
+            saver = AsyncCouchbaseSaver(
+                cluster,
+                bucket_name,
+                scope_name,
+                checkpoints_collection_name,
+                checkpoint_writes_collection_name,
+            )
+
+            await saver.create_collections()
             
             yield saver
         finally:
             if cluster:
                 await cluster.close()
 
+
     @classmethod
     @asynccontextmanager
     async def from_cluster(
@@ -98,9 +138,18 @@ async def from_cluster(
             AsyncCouchbaseSaver: An instance of the AsyncCouchbaseSaver
         """        
         
-        saver = AsyncCouchbaseSaver(cluster, bucket_name, scope_name, checkpoints_collection_name, checkpoint_writes_collection_name)
-        cls.bucket = cluster.bucket(bucket_name)
-        await cls.bucket.on_connect()
+        bucket = cluster.bucket(bucket_name)
+        await bucket.on_connect()
+
+        saver = AsyncCouchbaseSaver(
+            cluster,
+            bucket_name,
+            scope_name,
+            checkpoints_collection_name,
+            checkpoint_writes_collection_name,
+        )
+
+        await saver.create_collections()
         
         yield saver
 
@@ -149,7 +198,7 @@ async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
             async for write_doc in serialized_writes_result:
                 checkpoint_writes = write_doc.get(self.checkpoint_writes_collection_name, {})
                 if "task_id" not in checkpoint_writes:
-                    print("Error: 'task_id' is not present in checkpoint_writes")
+                    logger.warning("'task_id' is not present in checkpoint_writes")
                 else:
                     pending_writes.append(
                         (
@@ -294,6 +343,8 @@ async def aput(
 
         upsert_key = f"{thread_id}::{checkpoint_ns}::{checkpoint_id}"
 
+        # ensure bucket connected (idempotent)
+        await self.bucket.on_connect()
         collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
         await collection.upsert(upsert_key, (doc), UpsertOptions(timeout=timedelta(seconds=5)))
 
@@ -324,6 +375,7 @@ async def aput_writes(
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = config["configurable"]["checkpoint_id"]
 
+        await self.bucket.on_connect()
         collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
 
         for idx, (channel, value) in enumerate(writes):
diff --git a/langgraph_checkpointer_couchbase/couchbase_saver.py b/langgraph_checkpointer_couchbase/couchbase_saver.py
@@ -1,6 +1,7 @@
 from contextlib import contextmanager
 from datetime import timedelta
 from typing import Any, Dict, Iterator, Optional, Sequence, Tuple
+import logging
 
 from langchain_core.runnables import RunnableConfig
 from couchbase.cluster import Cluster
@@ -19,6 +20,8 @@
 )
 
 from .utils import _encode_binary, _decode_binary
+
+logger = logging.getLogger(__name__)
 class CouchbaseSaver(BaseCheckpointSaver):
     """A checkpoint saver that stores checkpoints in a Couchbase database.
     
@@ -118,7 +121,7 @@ def create_collections(self):
         except CollectionAlreadyExistsException as _:
             pass
         except Exception as e:
-            print(f"Error creating collections: {e}")
+            logger.exception("Error creating collections")
             raise e
         finally:
             self.checkpoints_collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
@@ -128,7 +131,7 @@ def create_collections(self):
         except CollectionAlreadyExistsException as _:
             pass
         except Exception as e:
-            print(f"Error creating collections: {e}")
+            logger.exception("Error creating collections")
             raise e
         finally:
             self.checkpoint_writes_collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
@@ -181,7 +184,7 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
             for write_doc in serialized_writes_result:
                 checkpoint_writes = write_doc.get(self.checkpoint_writes_collection_name, {})
                 if "task_id" not in checkpoint_writes:
-                    print("Error: 'task_id' is not present in checkpoint_writes")
+                    logger.warning("'task_id' is not present in checkpoint_writes")
                 else:
                     # Decode and deserialize value data
                     value_data = _decode_binary(checkpoint_writes["value"])
diff --git a/tests/agent_e2e_test.py b/tests/agent_e2e_test.py
@@ -4,7 +4,7 @@
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
 from langgraph.prebuilt import create_react_agent
-from langgraph_checkpoint_couchbase import CouchbaseSaver, AsyncCouchbaseSaver
+from langgraph_checkpointer_couchbase import CouchbaseSaver, AsyncCouchbaseSaver
 from dotenv import load_dotenv
 import os
 load_dotenv()
