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.
- langgraph_checkpointer_couchbase-1.0.1/.env.example +6 -0
- langgraph_checkpointer_couchbase-1.0.1/.github/workflows/release.yaml +45 -0
- langgraph_checkpointer_couchbase-1.0.1/.gitignore +13 -0
- langgraph_checkpointer_couchbase-1.0.1/LICENSE +9 -0
- langgraph_checkpointer_couchbase-1.0.1/PKG-INFO +179 -0
- langgraph_checkpointer_couchbase-1.0.1/README.md +152 -0
- langgraph_checkpointer_couchbase-1.0.1/langgraph_checkpointer_couchbase/__init__.py +4 -0
- langgraph_checkpointer_couchbase-1.0.1/langgraph_checkpointer_couchbase/async_cb_saver.py +346 -0
- langgraph_checkpointer_couchbase-1.0.1/langgraph_checkpointer_couchbase/couchbase_saver.py +390 -0
- langgraph_checkpointer_couchbase-1.0.1/langgraph_checkpointer_couchbase/utils.py +11 -0
- langgraph_checkpointer_couchbase-1.0.1/pyproject.toml +67 -0
- langgraph_checkpointer_couchbase-1.0.1/tests/__init__.py +0 -0
- langgraph_checkpointer_couchbase-1.0.1/tests/agent_e2e_test.py +70 -0
|
@@ -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,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
|
+
[](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
|
+
[](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,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
|
+
]
|
|
File without changes
|
|
@@ -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())
|