langgraph-checkpointer-couchbase 1.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,6 @@
1
+ OPENAI_API_KEY=
2
+ CB_CLUSTER=
3
+ CB_USERNAME=
4
+ CB_PASSWORD=
5
+ CB_BUCKET=
6
+ CB_SCOPE=
@@ -0,0 +1,45 @@
1
+ name: Publish Package to PyPI
2
+ on:
3
+ push:
4
+ branches:
5
+ - main
6
+ workflow_dispatch:
7
+
8
+ jobs:
9
+ publish:
10
+ runs-on: ubuntu-latest
11
+ permissions:
12
+ id-token: write
13
+ contents: write
14
+ steps:
15
+ - name: Checkout Repository
16
+ uses: actions/checkout@v3
17
+
18
+ - name: Set up Python
19
+ uses: actions/setup-python@v4
20
+ with:
21
+ python-version: "3.11"
22
+
23
+ - name: Install Dependencies
24
+ run: pip install -e .
25
+
26
+ - name: Install Hatch
27
+ run: pip install hatch
28
+
29
+ - name: Hatch build
30
+ run: hatch build
31
+
32
+ - name: Get hatch version
33
+ id: hatch_version_step
34
+ run: echo "VERSION=$(hatch version)" >> $GITHUB_OUTPUT
35
+
36
+ - name: Publish release distributions to PyPI
37
+ uses: pypa/gh-action-pypi-publish@release/v1
38
+
39
+ - name: Create GitHub Release
40
+ id: create_release
41
+ uses: softprops/action-gh-release@v2
42
+ with:
43
+ tag_name: v${{ steps.hatch_version_step.outputs.VERSION }}
44
+ draft: false
45
+ prerelease: false
@@ -0,0 +1,13 @@
1
+ .venv/
2
+ .env
3
+
4
+ __pycache__/
5
+ dist/
6
+ build/
7
+
8
+ *.pyc
9
+ *.pyo
10
+ *.pyd
11
+ *.pyw
12
+ *.pyz
13
+
@@ -0,0 +1,9 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
6
+
7
+ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
8
+
9
+ THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
@@ -0,0 +1,179 @@
1
+ Metadata-Version: 2.4
2
+ Name: langgraph-checkpointer-couchbase
3
+ Version: 1.0.1
4
+ Project-URL: Documentation, https://github.com/Lokesh Goel/langgraph-checkpointer-couchbase#readme
5
+ Project-URL: Issues, https://github.com/Lokesh Goel/langgraph-checkpointer-couchbase/issues
6
+ Project-URL: Source, https://github.com/Lokesh Goel/langgraph-checkpointer-couchbase
7
+ Author-email: Lokesh Goel <lokesh.goel@couchbase.com>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: checkpointer,couchbase,langchain,langgraph,persistence
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Programming Language :: Python
13
+ Classifier: Programming Language :: Python :: 3.8
14
+ Classifier: Programming Language :: Python :: 3.9
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: Implementation :: CPython
19
+ Classifier: Programming Language :: Python :: Implementation :: PyPy
20
+ Requires-Python: >=3.8
21
+ Requires-Dist: couchbase>=4.3.5
22
+ Requires-Dist: langchain-openai>=0.3.11
23
+ Requires-Dist: langgraph>=0.3.22
24
+ Requires-Dist: pydantic>=2.11.1
25
+ Requires-Dist: typing-extensions>=4.13.0
26
+ Description-Content-Type: text/markdown
27
+
28
+ # LangGraph Checkpoint Couchbase
29
+
30
+ A Couchbase implementation of the LangGraph `CheckpointSaver` interface that enables persisting agent state and conversation history in a Couchbase database.
31
+
32
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
33
+
34
+ ## Overview
35
+
36
+ This package provides a seamless way to persist LangGraph agent states in Couchbase, enabling:
37
+ - State persistence across application restarts
38
+ - Retrieval of historical conversation steps
39
+ - Continued conversations from previous checkpoints
40
+ - Both synchronous and asynchronous interfaces
41
+
42
+ ## Installation
43
+
44
+ ```bash
45
+ pip install langgraph-checkpointer-couchbase
46
+ ```
47
+
48
+ ## Requirements
49
+
50
+ - Python 3.8+
51
+ - Couchbase Server (7.0+ recommended)
52
+ - LangGraph 0.3.22+
53
+ - LangChain OpenAI 0.3.11+
54
+
55
+ ## Prerequisites
56
+
57
+ - A running Couchbase cluster
58
+ - A bucket created for storing checkpoints
59
+ - Appropriate credentials with read/write access
60
+
61
+ ## Quick Start
62
+
63
+ First, set up your agent tools and model:
64
+
65
+ ```python
66
+ from typing import Literal
67
+ from langchain_openai import ChatOpenAI
68
+
69
+ @tool
70
+ def get_weather(city: Literal["nyc", "sf"]):
71
+ """Use this to get weather information."""
72
+ if city == "nyc":
73
+ return "It might be cloudy in nyc"
74
+ elif city == "sf":
75
+ return "It's always sunny in sf"
76
+ else:
77
+ raise AssertionError("Unknown city")
78
+
79
+
80
+ tools = [get_weather]
81
+ model = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
82
+ ```
83
+
84
+ ### Synchronous Usage
85
+
86
+ ```python
87
+ import os
88
+ from langgraph_checkpointer_couchbase import CouchbaseSaver
89
+ from langgraph.graph import create_react_agent
90
+
91
+ with CouchbaseSaver.from_conn_info(
92
+ cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
93
+ cb_username=os.getenv("CB_USERNAME") or "Administrator",
94
+ cb_password=os.getenv("CB_PASSWORD") or "password",
95
+ bucket_name=os.getenv("CB_BUCKET") or "test",
96
+ scope_name=os.getenv("CB_SCOPE") or "langgraph",
97
+ ) as checkpointer:
98
+ # Create the agent with checkpointing
99
+ graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
100
+
101
+ # Configure with a unique thread ID
102
+ config = {"configurable": {"thread_id": "1"}}
103
+
104
+ # Run the agent
105
+ res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)
106
+
107
+ # Retrieve checkpoints
108
+ latest_checkpoint = checkpointer.get(config)
109
+ latest_checkpoint_tuple = checkpointer.get_tuple(config)
110
+ checkpoint_tuples = list(checkpointer.list(config))
111
+
112
+ print(latest_checkpoint)
113
+ print(latest_checkpoint_tuple)
114
+ print(checkpoint_tuples)
115
+ ```
116
+
117
+ ### Asynchronous Usage
118
+
119
+ ```python
120
+ import os
121
+ from langgraph_checkpointer_couchbase import AsyncCouchbaseSaver
122
+ from langgraph.graph import create_react_agent
123
+
124
+ async with AsyncCouchbaseSaver.from_conn_info(
125
+ cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
126
+ cb_username=os.getenv("CB_USERNAME") or "Administrator",
127
+ cb_password=os.getenv("CB_PASSWORD") or "password",
128
+ bucket_name=os.getenv("CB_BUCKET") or "test",
129
+ scope_name=os.getenv("CB_SCOPE") or "langgraph",
130
+ ) as checkpointer:
131
+ # Create the agent with checkpointing
132
+ graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
133
+
134
+ # Configure with a unique thread ID
135
+ config = {"configurable": {"thread_id": "2"}}
136
+
137
+ # Run the agent asynchronously
138
+ res = await graph.ainvoke(
139
+ {"messages": [("human", "what's the weather in nyc")]}, config
140
+ )
141
+
142
+ # Retrieve checkpoints asynchronously
143
+ latest_checkpoint = await checkpointer.aget(config)
144
+ latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
145
+ checkpoint_tuples = [c async for c in checkpointer.alist(config)]
146
+
147
+ print(latest_checkpoint)
148
+ print(latest_checkpoint_tuple)
149
+ print(checkpoint_tuples)
150
+ ```
151
+
152
+ ## Configuration Options
153
+
154
+ | Parameter | Description | Default |
155
+ |-----------|-------------|---------|
156
+ | CB_CLUSTER | Couchbase connection string | couchbase://localhost |
157
+ | CB_USERNAME | Username for Couchbase | Administrator |
158
+ | CB_PASSWORD | Password for Couchbase | password |
159
+ | CB_BUCKET | Bucket to store checkpoints | test |
160
+ | CB_SCOPE | Scope within bucket | langgraph |
161
+
162
+ ## License
163
+
164
+ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
165
+
166
+ ## Contributing
167
+
168
+ Contributions are welcome! Please feel free to submit a Pull Request.
169
+ ---
170
+
171
+ # 📢 Support Policy
172
+
173
+ We truly appreciate your interest in this project!
174
+ This project is **community-maintained**, which means it's **not officially supported** by our support team.
175
+
176
+ If you need help, have found a bug, or want to contribute improvements, the best place to do that is right here — by [opening a GitHub issue](https://github.com/Couchbase-Ecosystem/langgraph-checkpointer-couchbase/issues).
177
+ Our support portal is unable to assist with requests related to this project, so we kindly ask that all inquiries stay within GitHub.
178
+
179
+ Your collaboration helps us all move forward together — thank you!
@@ -0,0 +1,152 @@
1
+ # LangGraph Checkpoint Couchbase
2
+
3
+ A Couchbase implementation of the LangGraph `CheckpointSaver` interface that enables persisting agent state and conversation history in a Couchbase database.
4
+
5
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
6
+
7
+ ## Overview
8
+
9
+ This package provides a seamless way to persist LangGraph agent states in Couchbase, enabling:
10
+ - State persistence across application restarts
11
+ - Retrieval of historical conversation steps
12
+ - Continued conversations from previous checkpoints
13
+ - Both synchronous and asynchronous interfaces
14
+
15
+ ## Installation
16
+
17
+ ```bash
18
+ pip install langgraph-checkpointer-couchbase
19
+ ```
20
+
21
+ ## Requirements
22
+
23
+ - Python 3.8+
24
+ - Couchbase Server (7.0+ recommended)
25
+ - LangGraph 0.3.22+
26
+ - LangChain OpenAI 0.3.11+
27
+
28
+ ## Prerequisites
29
+
30
+ - A running Couchbase cluster
31
+ - A bucket created for storing checkpoints
32
+ - Appropriate credentials with read/write access
33
+
34
+ ## Quick Start
35
+
36
+ First, set up your agent tools and model:
37
+
38
+ ```python
39
+ from typing import Literal
40
+ from langchain_openai import ChatOpenAI
41
+
42
+ @tool
43
+ def get_weather(city: Literal["nyc", "sf"]):
44
+ """Use this to get weather information."""
45
+ if city == "nyc":
46
+ return "It might be cloudy in nyc"
47
+ elif city == "sf":
48
+ return "It's always sunny in sf"
49
+ else:
50
+ raise AssertionError("Unknown city")
51
+
52
+
53
+ tools = [get_weather]
54
+ model = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
55
+ ```
56
+
57
+ ### Synchronous Usage
58
+
59
+ ```python
60
+ import os
61
+ from langgraph_checkpointer_couchbase import CouchbaseSaver
62
+ from langgraph.graph import create_react_agent
63
+
64
+ with CouchbaseSaver.from_conn_info(
65
+ cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
66
+ cb_username=os.getenv("CB_USERNAME") or "Administrator",
67
+ cb_password=os.getenv("CB_PASSWORD") or "password",
68
+ bucket_name=os.getenv("CB_BUCKET") or "test",
69
+ scope_name=os.getenv("CB_SCOPE") or "langgraph",
70
+ ) as checkpointer:
71
+ # Create the agent with checkpointing
72
+ graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
73
+
74
+ # Configure with a unique thread ID
75
+ config = {"configurable": {"thread_id": "1"}}
76
+
77
+ # Run the agent
78
+ res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)
79
+
80
+ # Retrieve checkpoints
81
+ latest_checkpoint = checkpointer.get(config)
82
+ latest_checkpoint_tuple = checkpointer.get_tuple(config)
83
+ checkpoint_tuples = list(checkpointer.list(config))
84
+
85
+ print(latest_checkpoint)
86
+ print(latest_checkpoint_tuple)
87
+ print(checkpoint_tuples)
88
+ ```
89
+
90
+ ### Asynchronous Usage
91
+
92
+ ```python
93
+ import os
94
+ from langgraph_checkpointer_couchbase import AsyncCouchbaseSaver
95
+ from langgraph.graph import create_react_agent
96
+
97
+ async with AsyncCouchbaseSaver.from_conn_info(
98
+ cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
99
+ cb_username=os.getenv("CB_USERNAME") or "Administrator",
100
+ cb_password=os.getenv("CB_PASSWORD") or "password",
101
+ bucket_name=os.getenv("CB_BUCKET") or "test",
102
+ scope_name=os.getenv("CB_SCOPE") or "langgraph",
103
+ ) as checkpointer:
104
+ # Create the agent with checkpointing
105
+ graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
106
+
107
+ # Configure with a unique thread ID
108
+ config = {"configurable": {"thread_id": "2"}}
109
+
110
+ # Run the agent asynchronously
111
+ res = await graph.ainvoke(
112
+ {"messages": [("human", "what's the weather in nyc")]}, config
113
+ )
114
+
115
+ # Retrieve checkpoints asynchronously
116
+ latest_checkpoint = await checkpointer.aget(config)
117
+ latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
118
+ checkpoint_tuples = [c async for c in checkpointer.alist(config)]
119
+
120
+ print(latest_checkpoint)
121
+ print(latest_checkpoint_tuple)
122
+ print(checkpoint_tuples)
123
+ ```
124
+
125
+ ## Configuration Options
126
+
127
+ | Parameter | Description | Default |
128
+ |-----------|-------------|---------|
129
+ | CB_CLUSTER | Couchbase connection string | couchbase://localhost |
130
+ | CB_USERNAME | Username for Couchbase | Administrator |
131
+ | CB_PASSWORD | Password for Couchbase | password |
132
+ | CB_BUCKET | Bucket to store checkpoints | test |
133
+ | CB_SCOPE | Scope within bucket | langgraph |
134
+
135
+ ## License
136
+
137
+ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
138
+
139
+ ## Contributing
140
+
141
+ Contributions are welcome! Please feel free to submit a Pull Request.
142
+ ---
143
+
144
+ # 📢 Support Policy
145
+
146
+ We truly appreciate your interest in this project!
147
+ This project is **community-maintained**, which means it's **not officially supported** by our support team.
148
+
149
+ If you need help, have found a bug, or want to contribute improvements, the best place to do that is right here — by [opening a GitHub issue](https://github.com/Couchbase-Ecosystem/langgraph-checkpointer-couchbase/issues).
150
+ Our support portal is unable to assist with requests related to this project, so we kindly ask that all inquiries stay within GitHub.
151
+
152
+ Your collaboration helps us all move forward together — thank you!
@@ -0,0 +1,4 @@
1
+ from .async_cb_saver import AsyncCouchbaseSaver
2
+ from .couchbase_saver import CouchbaseSaver
3
+
4
+ __all__ = ["CouchbaseSaver", "AsyncCouchbaseSaver"]
@@ -0,0 +1,346 @@
1
+ from contextlib import asynccontextmanager
2
+ from datetime import timedelta
3
+ from typing import Any, AsyncIterator, Dict, Optional, Sequence, Tuple
4
+
5
+ from langchain_core.runnables import RunnableConfig
6
+ from acouchbase.cluster import Cluster as ACluster
7
+ from acouchbase.bucket import Bucket as ABucket
8
+ from couchbase.auth import PasswordAuthenticator
9
+ from couchbase.options import ClusterOptions, QueryOptions, UpsertOptions
10
+
11
+ from langgraph.checkpoint.base import (
12
+ BaseCheckpointSaver,
13
+ ChannelVersions,
14
+ Checkpoint,
15
+ CheckpointMetadata,
16
+ CheckpointTuple,
17
+ get_checkpoint_id,
18
+ )
19
+ from .utils import _encode_binary, _decode_binary
20
+
21
+ class AsyncCouchbaseSaver(BaseCheckpointSaver):
22
+ """A checkpoint saver that stores checkpoints in a Couchbase database."""
23
+
24
+ cluster: ACluster
25
+ bucket: ABucket
26
+ def __init__(
27
+ self,
28
+ cluster: ACluster,
29
+ bucket_name: str,
30
+ scope_name: str,
31
+ checkpoints_collection_name: str = "checkpoints",
32
+ checkpoint_writes_collection_name: str = "checkpoint_writes"
33
+ ) -> None:
34
+ super().__init__()
35
+ self.cluster = cluster
36
+ self.bucket_name = bucket_name
37
+ self.scope_name = scope_name
38
+ self.checkpoints_collection_name = checkpoints_collection_name
39
+ self.checkpoint_writes_collection_name = checkpoint_writes_collection_name
40
+
41
+ @classmethod
42
+ @asynccontextmanager
43
+ async def from_conn_info(
44
+ cls, *, cb_conn_str :str, cb_username: str, cb_password: str, bucket_name: str, scope_name: str, checkpoints_collection_name: str = "checkpoints", checkpoint_writes_collection_name: str = "checkpoint_writes"
45
+ ) -> AsyncIterator["AsyncCouchbaseSaver"]:
46
+ """Create a AsyncCouchbaseSaver from connection information.
47
+
48
+ This allows reusing an existing cluster connection rather than creating a new one.
49
+
50
+ Args:
51
+ cb_conn_str: Connection string for the Couchbase cluster
52
+ cb_username: Username for the Couchbase cluster
53
+ cb_password: Password for the Couchbase cluster
54
+ bucket_name: Name of the bucket to use
55
+ scope_name: Name of the scope within the bucket
56
+ checkpoints_collection_name: Name of the collection to store checkpoints
57
+ checkpoint_writes_collection_name: Name of the collection to store checkpoint writes
58
+
59
+ Yields:
60
+ AsyncCouchbaseSaver: An instance of the AsyncCouchbaseSaver
61
+ """
62
+ cluster = None
63
+ try:
64
+ auth = PasswordAuthenticator(cb_username, cb_password)
65
+ options = ClusterOptions(auth)
66
+ cluster = await ACluster.connect(cb_conn_str, options)
67
+
68
+ cls.cluster = cluster
69
+ cls.bucket_name = bucket_name
70
+ cls.scope_name = scope_name
71
+
72
+ saver = AsyncCouchbaseSaver(cluster, bucket_name, scope_name, checkpoints_collection_name, checkpoint_writes_collection_name)
73
+ cls.bucket = cluster.bucket(bucket_name)
74
+ await cls.bucket.on_connect()
75
+
76
+ yield saver
77
+ finally:
78
+ if cluster:
79
+ await cluster.close()
80
+
81
+ @classmethod
82
+ @asynccontextmanager
83
+ async def from_cluster(
84
+ cls, *, cluster: ACluster, bucket_name: str, scope_name: str, checkpoints_collection_name: str = "checkpoints", checkpoint_writes_collection_name: str = "checkpoint_writes"
85
+ ) -> AsyncIterator["AsyncCouchbaseSaver"]:
86
+ """Create a AsyncCouchbaseSaver from an existing cluster connection.
87
+
88
+ This allows reusing an existing cluster connection rather than creating a new one.
89
+
90
+ Args:
91
+ cluster: An existing Couchbase Cluster connection
92
+ bucket_name: Name of the bucket to use
93
+ scope_name: Name of the scope within the bucket
94
+ checkpoints_collection_name: Name of the collection to store checkpoints
95
+ checkpoint_writes_collection_name: Name of the collection to store checkpoint writes
96
+
97
+ Yields:
98
+ AsyncCouchbaseSaver: An instance of the AsyncCouchbaseSaver
99
+ """
100
+
101
+ saver = AsyncCouchbaseSaver(cluster, bucket_name, scope_name, checkpoints_collection_name, checkpoint_writes_collection_name)
102
+ cls.bucket = cluster.bucket(bucket_name)
103
+ await cls.bucket.on_connect()
104
+
105
+ yield saver
106
+
107
+ async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
108
+ """Get a checkpoint tuple from the database asynchronously.
109
+
110
+ This method retrieves a checkpoint tuple from the Couchbase database based on the
111
+ provided config. If the config contains a "checkpoint_id" key, the checkpoint with
112
+ the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
113
+ for the given thread ID is retrieved.
114
+
115
+ Args:
116
+ config (RunnableConfig): The config to use for retrieving the checkpoint.
117
+
118
+ Returns:
119
+ Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
120
+ """
121
+ thread_id = config["configurable"]["thread_id"]
122
+ checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
123
+ checkpoint_id = get_checkpoint_id(config)
124
+
125
+ if checkpoint_id:
126
+ query = f'SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoints_collection_name}` WHERE thread_id = $1 AND checkpoint_ns = $2 AND checkpoint_id = $3 ORDER BY checkpoint_id DESC LIMIT 1'
127
+ query_params = [thread_id, checkpoint_ns, checkpoint_id]
128
+ else:
129
+ query = f'SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoints_collection_name}` WHERE thread_id = $1 AND checkpoint_ns = $2 ORDER BY checkpoint_id DESC LIMIT 1'
130
+ query_params = [thread_id, checkpoint_ns]
131
+
132
+ print(query)
133
+ result = self.cluster.query(query, QueryOptions(positional_parameters=query_params))
134
+
135
+ async for row in result:
136
+ doc = row[self.checkpoints_collection_name]
137
+ config_values = {
138
+ "thread_id": thread_id,
139
+ "checkpoint_ns": checkpoint_ns,
140
+ "checkpoint_id": doc["checkpoint_id"],
141
+ }
142
+ checkpoint = self.serde.loads_typed((doc["type"], _decode_binary(doc["checkpoint"])))
143
+
144
+ serialized_writes_query = f'SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoint_writes_collection_name}` WHERE thread_id = $1 AND checkpoint_ns = $2 AND checkpoint_id = $3'
145
+ serialized_writes_params = [thread_id, checkpoint_ns, doc["checkpoint_id"] or ""]
146
+
147
+ serialized_writes_result = self.cluster.query(serialized_writes_query, QueryOptions(positional_parameters=serialized_writes_params))
148
+
149
+ pending_writes = []
150
+ async for write_doc in serialized_writes_result:
151
+ print(f"write_doc: {write_doc}") # Debugging statement to log the contents of write_doc
152
+ checkpoint_writes = write_doc.get(self.checkpoint_writes_collection_name, {})
153
+ if "task_id" not in checkpoint_writes:
154
+ print("Error: 'task_id' is not present in checkpoint_writes")
155
+ else:
156
+ pending_writes.append(
157
+ (
158
+ checkpoint_writes["task_id"],
159
+ checkpoint_writes["channel"],
160
+ self.serde.loads_typed((checkpoint_writes["type"], _decode_binary(checkpoint_writes["value"]))),
161
+ )
162
+ )
163
+
164
+ return CheckpointTuple(
165
+ {"configurable": config_values},
166
+ checkpoint,
167
+ self.serde.loads(_decode_binary(doc["metadata"])),
168
+ (
169
+ {
170
+ "configurable": {
171
+ "thread_id": thread_id,
172
+ "checkpoint_ns": checkpoint_ns,
173
+ "checkpoint_id": doc["parent_checkpoint_id"],
174
+ }
175
+ }
176
+ if doc.get("parent_checkpoint_id")
177
+ else None
178
+ ),
179
+ pending_writes,
180
+ )
181
+
182
+ async def alist(
183
+ self,
184
+ config: Optional[RunnableConfig],
185
+ *,
186
+ filter: Optional[Dict[str, Any]] = None,
187
+ before: Optional[RunnableConfig] = None,
188
+ limit: Optional[int] = None,
189
+ ) -> AsyncIterator[CheckpointTuple]:
190
+ """List checkpoints from the database asynchronously.
191
+
192
+ This method retrieves a list of checkpoint tuples from the Couchbase database based
193
+ on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
194
+
195
+ Args:
196
+ config (Optional[RunnableConfig]): The config to use for listing the checkpoints.
197
+ filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
198
+ before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
199
+ limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
200
+
201
+ Yields:
202
+ AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
203
+ """
204
+
205
+ query = f"SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoints_collection_name}` WHERE 1=1"
206
+ query_params = []
207
+
208
+ if config is not None:
209
+ query += " AND thread_id = $1 AND checkpoint_ns = $2"
210
+ query_params.extend([config["configurable"]["thread_id"], config["configurable"].get("checkpoint_ns", "")])
211
+
212
+ if filter:
213
+ for key, value in filter.items():
214
+ query += f" AND metadata.{key} = ${len(query_params) + 1}"
215
+ query_params.append(value)
216
+
217
+ if before is not None:
218
+ query += f" AND checkpoint_id < ${len(query_params) + 1}"
219
+ query_params.append(before["configurable"]["checkpoint_id"])
220
+
221
+ query += " ORDER BY checkpoint_id DESC"
222
+
223
+ if limit is not None:
224
+ query += f" LIMIT {limit}"
225
+
226
+ result = self.cluster.query(query, QueryOptions(positional_parameters=query_params))
227
+
228
+ async for row in result:
229
+ doc = row[self.checkpoints_collection_name]
230
+ checkpoint_data = _decode_binary(doc["checkpoint"])
231
+ checkpoint = self.serde.loads_typed((doc["type"], checkpoint_data))
232
+ yield CheckpointTuple(
233
+ {
234
+ "configurable": {
235
+ "thread_id": doc["thread_id"],
236
+ "checkpoint_ns": doc["checkpoint_ns"],
237
+ "checkpoint_id": doc["checkpoint_id"],
238
+ }
239
+ },
240
+ checkpoint,
241
+ self.serde.loads(_decode_binary(doc["metadata"])),
242
+ (
243
+ {
244
+ "configurable": {
245
+ "thread_id": doc["thread_id"],
246
+ "checkpoint_ns": doc["checkpoint_ns"],
247
+ "checkpoint_id": doc["parent_checkpoint_id"],
248
+ }
249
+ }
250
+ if doc.get("parent_checkpoint_id")
251
+ else None
252
+ ),
253
+ )
254
+
255
+ async def aput(
256
+ self,
257
+ config: RunnableConfig,
258
+ checkpoint: Checkpoint,
259
+ metadata: CheckpointMetadata,
260
+ new_versions: ChannelVersions,
261
+ ) -> RunnableConfig:
262
+ """Save a checkpoint to the database asynchronously.
263
+
264
+ This method saves a checkpoint to the Couchbase database. The checkpoint is associated
265
+ with the provided config and its parent config (if any).
266
+
267
+ Args:
268
+ config (RunnableConfig): The config to associate with the checkpoint.
269
+ checkpoint (Checkpoint): The checkpoint to save.
270
+ metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
271
+ new_versions (ChannelVersions): New channel versions as of this write.
272
+
273
+ Returns:
274
+ RunnableConfig: Updated configuration after storing the checkpoint.
275
+ """
276
+ thread_id = config["configurable"]["thread_id"]
277
+ checkpoint_ns = config["configurable"]["checkpoint_ns"]
278
+ checkpoint_id = checkpoint["id"]
279
+ type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
280
+ if serialized_checkpoint:
281
+ serialized_checkpoint = _encode_binary(serialized_checkpoint)
282
+
283
+ metadata = self.serde.dumps(metadata)
284
+ if metadata:
285
+ metadata = _encode_binary(metadata)
286
+
287
+ doc = {
288
+ "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
289
+ "type": type_,
290
+ "checkpoint": serialized_checkpoint,
291
+ "metadata": metadata,
292
+ "thread_id" : thread_id,
293
+ "checkpoint_ns": checkpoint_ns,
294
+ "checkpoint_id": checkpoint_id,
295
+ }
296
+ # print(json.dumps(doc))
297
+ upsert_key = f"{thread_id}::{checkpoint_ns}::{checkpoint_id}"
298
+
299
+ collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
300
+ await collection.upsert(upsert_key, (doc), UpsertOptions(timeout=timedelta(seconds=5)))
301
+
302
+ return {
303
+ "configurable": {
304
+ "thread_id": thread_id,
305
+ "checkpoint_ns": checkpoint_ns,
306
+ "checkpoint_id": checkpoint_id,
307
+ }
308
+ }
309
+
310
+ async def aput_writes(
311
+ self,
312
+ config: RunnableConfig,
313
+ writes: Sequence[Tuple[str, Any]],
314
+ task_id: str,
315
+ ) -> None:
316
+ """Store intermediate writes linked to a checkpoint asynchronously.
317
+
318
+ This method saves intermediate writes associated with a checkpoint to the Couchbase database.
319
+
320
+ Args:
321
+ config (RunnableConfig): Configuration of the related checkpoint.
322
+ writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
323
+ task_id (str): Identifier for the task creating the writes.
324
+ """
325
+ thread_id = config["configurable"]["thread_id"]
326
+ checkpoint_ns = config["configurable"]["checkpoint_ns"]
327
+ checkpoint_id = config["configurable"]["checkpoint_id"]
328
+
329
+ collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
330
+
331
+ for idx, (channel, value) in enumerate(writes):
332
+ upsert_key = f"{thread_id}::{checkpoint_ns}::{checkpoint_id}::{task_id}::{idx}"
333
+ type_, serialized_value = self.serde.dumps_typed(value)
334
+ if serialized_value:
335
+ serialized_value = _encode_binary(serialized_value)
336
+ doc = {
337
+ "thread_id": thread_id,
338
+ "checkpoint_ns": checkpoint_ns,
339
+ "checkpoint_id": checkpoint_id,
340
+ "task_id": task_id,
341
+ "idx": idx,
342
+ "channel": channel,
343
+ "type": type_,
344
+ "value": serialized_value,
345
+ }
346
+ await collection.upsert(upsert_key, (doc), UpsertOptions(timeout=timedelta(seconds=5)))
@@ -0,0 +1,390 @@
1
+ from contextlib import contextmanager
2
+ from datetime import timedelta
3
+ from typing import Any, Dict, Iterator, Optional, Sequence, Tuple
4
+
5
+ from langchain_core.runnables import RunnableConfig
6
+ from couchbase.cluster import Cluster
7
+ from couchbase.bucket import Bucket
8
+ from couchbase.auth import PasswordAuthenticator
9
+ from couchbase.options import ClusterOptions, QueryOptions, UpsertOptions
10
+ from couchbase.management.collections import CollectionSpec
11
+ from couchbase.exceptions import CollectionAlreadyExistsException
12
+
13
+ from langgraph.checkpoint.base import (
14
+ BaseCheckpointSaver,
15
+ ChannelVersions,
16
+ Checkpoint,
17
+ CheckpointMetadata,
18
+ CheckpointTuple,
19
+ get_checkpoint_id,
20
+ )
21
+
22
+ from .utils import _encode_binary, _decode_binary
23
+ class CouchbaseSaver(BaseCheckpointSaver):
24
+ """A checkpoint saver that stores checkpoints in a Couchbase database.
25
+
26
+ Initliaztion should be done using from_conn_info or from_cluster.
27
+ """
28
+
29
+ cluster: Cluster
30
+ bucket: Bucket
31
+ def __init__(
32
+ self,
33
+ cluster: Cluster,
34
+ bucket_name: str,
35
+ scope_name: str,
36
+ checkpoints_collection_name: str,
37
+ checkpoint_writes_collection_name: str,
38
+ ) -> None:
39
+ super().__init__()
40
+ self.cluster = cluster
41
+ self.bucket_name = bucket_name
42
+ self.scope_name = scope_name
43
+ self.bucket = self.cluster.bucket(bucket_name)
44
+ self.scope = self.bucket.scope(scope_name)
45
+ self.checkpoints_collection_name = checkpoints_collection_name
46
+ self.checkpoint_writes_collection_name = checkpoint_writes_collection_name
47
+
48
+ self.create_collections()
49
+
50
+ @classmethod
51
+ @contextmanager
52
+ def from_conn_info(
53
+ cls, *, cb_conn_str :str, cb_username: str, cb_password: str, bucket_name: str, scope_name: str, checkpoints_collection_name: str = "checkpoints", checkpoint_writes_collection_name: str = "checkpoint_writes"
54
+ ) -> Iterator["CouchbaseSaver"]:
55
+ """Create a CouchbaseSaver from connection information.
56
+
57
+ Cluster connection is created and closed automatically.
58
+
59
+ Args:
60
+ cb_conn_str: Connection string for the Couchbase cluster
61
+ cb_username: Username for the Couchbase cluster
62
+ cb_password: Password for the Couchbase cluster
63
+ bucket_name: Name of the bucket to use
64
+ scope_name: Name of the scope within the bucket
65
+ checkpoints_collection_name: Name of the collection to store checkpoints
66
+ checkpoint_writes_collection_name: Name of the collection to store checkpoint writes
67
+
68
+ Yields:
69
+ CouchbaseSaver: An instance of the CouchbaseSaver
70
+ """
71
+
72
+ cluster = None
73
+ try:
74
+ # Connect to Couchbase Cluster
75
+ auth = PasswordAuthenticator(cb_username, cb_password)
76
+ options = ClusterOptions(auth)
77
+ cluster = Cluster(cb_conn_str, options)
78
+ cluster.wait_until_ready(timedelta(seconds=5))
79
+
80
+ cls.cluster = cluster
81
+ cls.bucket_name = bucket_name
82
+ cls.scope_name = scope_name
83
+
84
+ saver = CouchbaseSaver(cluster, bucket_name, scope_name, checkpoints_collection_name, checkpoint_writes_collection_name)
85
+
86
+ yield saver
87
+ finally:
88
+ if cluster:
89
+ cluster.close()
90
+
91
+ @classmethod
92
+ @contextmanager
93
+ def from_cluster(
94
+ cls, *, cluster: Cluster, bucket_name: str, scope_name: str, checkpoints_collection_name: str = "checkpoints", checkpoint_writes_collection_name: str = "checkpoint_writes"
95
+ ) -> Iterator["CouchbaseSaver"]:
96
+ """Create a CouchbaseSaver from an existing cluster connection.
97
+
98
+ This allows reusing an existing cluster connection rather than creating a new one.
99
+
100
+ Args:
101
+ cluster: An existing Couchbase Cluster connection
102
+ bucket_name: Name of the bucket to use
103
+ scope_name: Name of the scope within the bucket
104
+
105
+ Yields:
106
+ CouchbaseSaver: An instance of the CouchbaseSaver
107
+ """
108
+
109
+ saver = CouchbaseSaver(cluster, bucket_name, scope_name, checkpoints_collection_name, checkpoint_writes_collection_name)
110
+
111
+ yield saver
112
+
113
+ def create_collections(self):
114
+ """Create collections in the Couchbase bucket if they do not exist."""
115
+
116
+ collection_manager = self.bucket.collections()
117
+ try:
118
+ collection_spec = CollectionSpec(collection_name=self.checkpoints_collection_name, scope_name=self.scope_name)
119
+ collection_manager.create_collection(collection_spec)
120
+ except CollectionAlreadyExistsException as caee:
121
+ pass
122
+ except Exception as e:
123
+ print(f"Error creating collections: {e}")
124
+ raise e
125
+ finally:
126
+ self.checkpoints_collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
127
+
128
+ try:
129
+ collection_spec = CollectionSpec(collection_name=self.checkpoint_writes_collection_name, scope_name=self.scope_name)
130
+ collection_manager.create_collection(collection_spec)
131
+ except CollectionAlreadyExistsException as caee:
132
+ pass
133
+ except Exception as e:
134
+ print(f"Error creating collections: {e}")
135
+ raise e
136
+ finally:
137
+ self.checkpoint_writes_collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
138
+
139
+ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
140
+ """Get a checkpoint tuple from the database.
141
+
142
+ This method retrieves a checkpoint tuple from the Couchbase database based on the
143
+ provided config. If the config contains a "checkpoint_id" key, the checkpoint with
144
+ the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
145
+ for the given thread ID is retrieved.
146
+
147
+ Args:
148
+ config (RunnableConfig): The config to use for retrieving the checkpoint.
149
+
150
+ Returns:
151
+ Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
152
+ """
153
+
154
+ thread_id = config["configurable"]["thread_id"]
155
+ checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
156
+ checkpoint_id = get_checkpoint_id(config)
157
+
158
+ if checkpoint_id:
159
+ query = f'SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoints_collection_name}` WHERE thread_id = $1 AND checkpoint_ns = $2 AND checkpoint_id = $3 ORDER BY checkpoint_id DESC LIMIT 1'
160
+ query_params = [thread_id, checkpoint_ns, checkpoint_id]
161
+ else:
162
+ query = f'SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoints_collection_name}` WHERE thread_id = $1 AND checkpoint_ns = $2 ORDER BY checkpoint_id DESC LIMIT 1'
163
+ query_params = [thread_id, checkpoint_ns]
164
+
165
+ result = self.cluster.query(query, QueryOptions(positional_parameters=query_params))
166
+
167
+ for row in result:
168
+ doc = row[self.checkpoints_collection_name]
169
+ config_values = {
170
+ "thread_id": thread_id,
171
+ "checkpoint_ns": checkpoint_ns,
172
+ "checkpoint_id": doc["checkpoint_id"],
173
+ }
174
+
175
+ # Decode and deserialize checkpoint data
176
+ checkpoint_data = _decode_binary(doc["checkpoint"])
177
+ checkpoint = self.serde.loads_typed((doc["type"], checkpoint_data))
178
+
179
+ serialized_writes_query = f'SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoint_writes_collection_name}` WHERE thread_id = $1 AND checkpoint_ns = $2 AND checkpoint_id = $3'
180
+ serialized_writes_params = [thread_id, checkpoint_ns, doc["checkpoint_id"] or ""]
181
+ serialized_writes_result = self.cluster.query(serialized_writes_query, QueryOptions(positional_parameters=serialized_writes_params))
182
+
183
+ pending_writes = []
184
+ for write_doc in serialized_writes_result:
185
+ checkpoint_writes = write_doc.get(self.checkpoint_writes_collection_name, {})
186
+ if "task_id" not in checkpoint_writes:
187
+ print("Error: 'task_id' is not present in checkpoint_writes")
188
+ else:
189
+ # Decode and deserialize value data
190
+ value_data = _decode_binary(checkpoint_writes["value"])
191
+ pending_writes.append(
192
+ (
193
+ checkpoint_writes["task_id"],
194
+ checkpoint_writes["channel"],
195
+ self.serde.loads_typed((checkpoint_writes["type"], value_data)),
196
+ )
197
+ )
198
+
199
+ # Decode and deserialize metadata
200
+ metadata = _decode_binary(doc["metadata"])
201
+ metadata = self.serde.loads(metadata)
202
+
203
+ return CheckpointTuple(
204
+ {"configurable": config_values},
205
+ checkpoint,
206
+ metadata,
207
+ (
208
+ {
209
+ "configurable": {
210
+ "thread_id": thread_id,
211
+ "checkpoint_ns": checkpoint_ns,
212
+ "checkpoint_id": doc["parent_checkpoint_id"],
213
+ }
214
+ }
215
+ if doc.get("parent_checkpoint_id")
216
+ else None
217
+ ),
218
+ pending_writes,
219
+ )
220
+
221
+ def list(
222
+ self,
223
+ config: Optional[RunnableConfig],
224
+ *,
225
+ filter: Optional[Dict[str, Any]] = None,
226
+ before: Optional[RunnableConfig] = None,
227
+ limit: Optional[int] = None,
228
+ ) -> Iterator[CheckpointTuple]:
229
+ """List checkpoints from the database.
230
+
231
+ This method retrieves a list of checkpoint tuples from the Couchbase database based
232
+ on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
233
+
234
+ Args:
235
+ config (RunnableConfig): The config to use for listing the checkpoints.
236
+ filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
237
+ before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
238
+ limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
239
+
240
+ Yields:
241
+ Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
242
+ """
243
+
244
+ query = f"SELECT * FROM {self.bucket_name}.{self.scope_name}.`{self.checkpoints_collection_name}` WHERE 1=1"
245
+ query_params = []
246
+
247
+ if config is not None:
248
+ query += " AND thread_id = $1 AND checkpoint_ns = $2"
249
+ query_params.extend([config["configurable"]["thread_id"], config["configurable"].get("checkpoint_ns", "")])
250
+
251
+ if filter:
252
+ for key, value in filter.items():
253
+ query += f" AND metadata.{key} = ${len(query_params) + 1}"
254
+ query_params.append(value)
255
+
256
+ if before is not None:
257
+ query += f" AND checkpoint_id < ${len(query_params) + 1}"
258
+ query_params.append(before["configurable"]["checkpoint_id"])
259
+
260
+ query += " ORDER BY checkpoint_id DESC"
261
+
262
+ if limit is not None:
263
+ query += f" LIMIT {limit}"
264
+
265
+ result = self.cluster.query(query, QueryOptions(positional_parameters=query_params))
266
+
267
+ for row in result:
268
+ doc = row[self.checkpoints_collection_name]
269
+ checkpoint_data = _decode_binary(doc["checkpoint"])
270
+ checkpoint = self.serde.loads_typed((doc["type"], checkpoint_data))
271
+ yield CheckpointTuple(
272
+ {
273
+ "configurable": {
274
+ "thread_id": doc["thread_id"],
275
+ "checkpoint_ns": doc["checkpoint_ns"],
276
+ "checkpoint_id": doc["checkpoint_id"],
277
+ }
278
+ },
279
+ checkpoint,
280
+ self.serde.loads(_decode_binary(doc["metadata"])),
281
+ (
282
+ {
283
+ "configurable": {
284
+ "thread_id": doc["thread_id"],
285
+ "checkpoint_ns": doc["checkpoint_ns"],
286
+ "checkpoint_id": doc["parent_checkpoint_id"],
287
+ }
288
+ }
289
+ if doc.get("parent_checkpoint_id")
290
+ else None
291
+ ),
292
+ )
293
+
294
+ def put(
295
+ self,
296
+ config: RunnableConfig,
297
+ checkpoint: Checkpoint,
298
+ metadata: CheckpointMetadata,
299
+ new_versions: ChannelVersions,
300
+ ) -> RunnableConfig:
301
+ """Save a checkpoint to the database.
302
+
303
+ This method saves a checkpoint to the Couchbase database. The checkpoint is associated
304
+ with the provided config and its parent config (if any).
305
+
306
+ Args:
307
+ config (RunnableConfig): The config to associate with the checkpoint.
308
+ checkpoint (Checkpoint): The checkpoint to save.
309
+ metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
310
+ new_versions (ChannelVersions): New channel versions as of this write.
311
+
312
+ Returns:
313
+ RunnableConfig: Updated configuration after storing the checkpoint.
314
+ """
315
+ thread_id = config["configurable"]["thread_id"]
316
+ checkpoint_ns = config["configurable"]["checkpoint_ns"]
317
+ checkpoint_id = checkpoint["id"]
318
+
319
+ # Serialize and encode checkpoint data
320
+ type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
321
+ if serialized_checkpoint:
322
+ serialized_checkpoint = _encode_binary(serialized_checkpoint)
323
+
324
+ # Serialize and encode metadata
325
+ metadata_bytes = self.serde.dumps(metadata)
326
+ if metadata_bytes:
327
+ metadata = _encode_binary(metadata_bytes)
328
+
329
+ doc = {
330
+ "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
331
+ "type": type_,
332
+ "checkpoint": serialized_checkpoint,
333
+ "metadata": metadata,
334
+ "thread_id" : thread_id,
335
+ "checkpoint_ns": checkpoint_ns,
336
+ "checkpoint_id": checkpoint_id,
337
+ }
338
+ upsert_key = f"{thread_id}::{checkpoint_ns}::{checkpoint_id}"
339
+
340
+ collection = self.checkpoints_collection
341
+ collection.upsert(upsert_key, (doc), UpsertOptions(timeout=timedelta(seconds=5)))
342
+
343
+ return {
344
+ "configurable": {
345
+ "thread_id": thread_id,
346
+ "checkpoint_ns": checkpoint_ns,
347
+ "checkpoint_id": checkpoint_id,
348
+ }
349
+ }
350
+
351
+ def put_writes(
352
+ self,
353
+ config: RunnableConfig,
354
+ writes: Sequence[Tuple[str, Any]],
355
+ task_id: str,
356
+ ) -> None:
357
+ """Store intermediate writes linked to a checkpoint.
358
+
359
+ This method saves intermediate writes associated with a checkpoint to the Couchbase database.
360
+
361
+ Args:
362
+ config (RunnableConfig): Configuration of the related checkpoint.
363
+ writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
364
+ task_id (str): Identifier for the task creating the writes.
365
+ """
366
+ thread_id = config["configurable"]["thread_id"]
367
+ checkpoint_ns = config["configurable"]["checkpoint_ns"]
368
+ checkpoint_id = config["configurable"]["checkpoint_id"]
369
+
370
+ collection = self.checkpoint_writes_collection
371
+
372
+ for idx, (channel, value) in enumerate(writes):
373
+ upsert_key = f"{thread_id}::{checkpoint_ns}::{checkpoint_id}::{task_id}::{idx}"
374
+ type_, serialized_value = self.serde.dumps_typed(value)
375
+
376
+ # Encode the serialized value to base64 string
377
+ if serialized_value:
378
+ serialized_value = _encode_binary(serialized_value)
379
+
380
+ doc = {
381
+ "thread_id": thread_id,
382
+ "checkpoint_ns": checkpoint_ns,
383
+ "checkpoint_id": checkpoint_id,
384
+ "task_id": task_id,
385
+ "idx": idx,
386
+ "channel": channel,
387
+ "type": type_,
388
+ "value": serialized_value,
389
+ }
390
+ collection.upsert(upsert_key, (doc), UpsertOptions(timeout=timedelta(seconds=5)))
@@ -0,0 +1,11 @@
1
+ import base64
2
+
3
+ __all__ = ["_encode_binary", "_decode_binary"]
4
+
5
+ def _encode_binary(data: bytes) -> str:
6
+ """Encode binary data to base64 string."""
7
+ return base64.b64encode(data).decode('utf-8')
8
+
9
+ def _decode_binary(data: str) -> bytes:
10
+ """Decode base64 string to binary data."""
11
+ return base64.b64decode(data)
@@ -0,0 +1,67 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "langgraph-checkpointer-couchbase"
7
+ version = "1.0.1"
8
+ description = ''
9
+ readme = "README.md"
10
+ requires-python = ">=3.8"
11
+ license = "MIT"
12
+ keywords = ["langchain","langgraph", "couchbase", "checkpointer", "persistence"]
13
+ authors = [
14
+ { name = "Lokesh Goel", email = "lokesh.goel@couchbase.com" },
15
+ ]
16
+ classifiers = [
17
+ "Development Status :: 4 - Beta",
18
+ "Programming Language :: Python",
19
+ "Programming Language :: Python :: 3.8",
20
+ "Programming Language :: Python :: 3.9",
21
+ "Programming Language :: Python :: 3.10",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: 3.12",
24
+ "Programming Language :: Python :: Implementation :: CPython",
25
+ "Programming Language :: Python :: Implementation :: PyPy",
26
+ ]
27
+ dependencies = [
28
+ "couchbase>=4.3.5",
29
+ "langgraph>=0.3.22",
30
+ "langchain-openai>=0.3.11",
31
+ "pydantic>=2.11.1",
32
+ "typing_extensions>=4.13.0"
33
+ ]
34
+
35
+ [project.urls]
36
+ Documentation = "https://github.com/Lokesh Goel/langgraph-checkpointer-couchbase#readme"
37
+ Issues = "https://github.com/Lokesh Goel/langgraph-checkpointer-couchbase/issues"
38
+ Source = "https://github.com/Lokesh Goel/langgraph-checkpointer-couchbase"
39
+
40
+ [tool.hatch.version]
41
+ path = "src/langgraph_checkpointer_couchbase/__about__.py"
42
+
43
+ [tool.hatch.envs.types]
44
+ extra-dependencies = [
45
+ "mypy>=1.0.0",
46
+ ]
47
+ [tool.hatch.envs.types.scripts]
48
+ check = "mypy --install-types --non-interactive {args:src/langgraph_checkpointer_couchbase tests}"
49
+
50
+ [tool.coverage.run]
51
+ source_pkgs = ["langgraph_checkpointer_couchbase", "tests"]
52
+ branch = true
53
+ parallel = true
54
+ omit = [
55
+ "src/langgraph_checkpointer_couchbase/__about__.py",
56
+ ]
57
+
58
+ [tool.coverage.paths]
59
+ langgraph_checkpointer_couchbase = ["src/langgraph_checkpointer_couchbase", "*/langgraph-checkpointer-couchbase/src/langgraph_checkpointer_couchbase"]
60
+ tests = ["tests", "*/langgraph-checkpointer-couchbase/tests"]
61
+
62
+ [tool.coverage.report]
63
+ exclude_lines = [
64
+ "no cov",
65
+ "if __name__ == .__main__.:",
66
+ "if TYPE_CHECKING:",
67
+ ]
@@ -0,0 +1,70 @@
1
+ import asyncio
2
+ from typing import Literal
3
+ from langchain_core.runnables import ConfigurableField
4
+ from langchain_core.tools import tool
5
+ from langchain_openai import ChatOpenAI
6
+ from langgraph.prebuilt import create_react_agent
7
+ from langgraph_checkpoint_couchbase import CouchbaseSaver, AsyncCouchbaseSaver
8
+ from dotenv import load_dotenv
9
+ import os
10
+ load_dotenv()
11
+
12
+ @tool
13
+ def get_weather(city: Literal["nyc", "sf"]):
14
+ """Use this to get weather information."""
15
+ if city == "nyc":
16
+ return "It might be cloudy in nyc"
17
+ elif city == "sf":
18
+ return "It's always sunny in sf"
19
+ else:
20
+ raise AssertionError("Unknown city")
21
+
22
+
23
+ tools = [get_weather]
24
+ model = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
25
+
26
+ def syncTest():
27
+ with CouchbaseSaver.from_conn_info(
28
+ cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
29
+ cb_username=os.getenv("CB_USERNAME") or "Administrator",
30
+ cb_password=os.getenv("CB_PASSWORD") or "password",
31
+ bucket_name=os.getenv("CB_BUCKET") or "test",
32
+ scope_name=os.getenv("CB_SCOPE") or "langgraph",
33
+ ) as checkpointer:
34
+ graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
35
+ config = {"configurable": {"thread_id": "1"}}
36
+ res = graph.invoke({"messages": [("human", "what's the weather in sf")]}, config)
37
+
38
+ latest_checkpoint = checkpointer.get(config)
39
+ latest_checkpoint_tuple = checkpointer.get_tuple(config)
40
+ checkpoint_tuples = list(checkpointer.list(config))
41
+
42
+ print(latest_checkpoint)
43
+ print(latest_checkpoint_tuple)
44
+ print(checkpoint_tuples)
45
+
46
+ async def asyncTest():
47
+ async with AsyncCouchbaseSaver.from_conn_info(
48
+ cb_conn_str=os.getenv("CB_CLUSTER") or "couchbase://localhost",
49
+ cb_username=os.getenv("CB_USERNAME") or "Administrator",
50
+ cb_password=os.getenv("CB_PASSWORD") or "password",
51
+ bucket_name=os.getenv("CB_BUCKET") or "test",
52
+ scope_name=os.getenv("CB_SCOPE") or "langgraph",
53
+ ) as checkpointer:
54
+ graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
55
+ config = {"configurable": {"thread_id": "2"}}
56
+ res = await graph.ainvoke(
57
+ {"messages": [("human", "what's the weather in nyc")]}, config
58
+ )
59
+
60
+ latest_checkpoint = await checkpointer.aget(config)
61
+ latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
62
+ checkpoint_tuples = [c async for c in checkpointer.alist(config)]
63
+
64
+ print(latest_checkpoint)
65
+ print(latest_checkpoint_tuple)
66
+ print(checkpoint_tuples)
67
+
68
+ if __name__ == "__main__":
69
+ syncTest()
70
+ asyncio.run(asyncTest())