langgraph-checkpointer-couchbase 1.0.5__tar.gz → 1.0.7__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.5 → langgraph_checkpointer_couchbase-1.0.7}/.github/workflows/release.yaml +3 -7
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/PKG-INFO +57 -43
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/README.md +56 -42
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/langgraph_checkpointer_couchbase/async_cb_saver.py +63 -11
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/langgraph_checkpointer_couchbase/couchbase_saver.py +10 -7
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/pyproject.toml +1 -1
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/tests/agent_e2e_test.py +1 -1
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/.env.example +0 -0
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/.gitignore +0 -0
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/LICENSE +0 -0
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/langgraph_checkpointer_couchbase/__init__.py +0 -0
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/langgraph_checkpointer_couchbase/utils.py +0 -0
- {langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/tests/__init__.py +0 -0
|
@@ -1,8 +1,8 @@
|
|
|
1
1
|
name: Publish Package to PyPI
|
|
2
2
|
on:
|
|
3
3
|
push:
|
|
4
|
-
|
|
5
|
-
-
|
|
4
|
+
tags:
|
|
5
|
+
- "v*"
|
|
6
6
|
workflow_dispatch:
|
|
7
7
|
|
|
8
8
|
jobs:
|
|
@@ -29,10 +29,6 @@ jobs:
|
|
|
29
29
|
- name: Hatch build
|
|
30
30
|
run: hatch build
|
|
31
31
|
|
|
32
|
-
- name: Get hatch version
|
|
33
|
-
id: hatch_version_step
|
|
34
|
-
run: echo "VERSION=$(hatch version)" >> $GITHUB_OUTPUT
|
|
35
|
-
|
|
36
32
|
- name: Publish release distributions to PyPI
|
|
37
33
|
uses: pypa/gh-action-pypi-publish@release/v1
|
|
38
34
|
|
|
@@ -40,6 +36,6 @@ jobs:
|
|
|
40
36
|
id: create_release
|
|
41
37
|
uses: softprops/action-gh-release@v2
|
|
42
38
|
with:
|
|
43
|
-
tag_name:
|
|
39
|
+
tag_name: ${{ github.ref_name }}
|
|
44
40
|
draft: false
|
|
45
41
|
prerelease: false
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: langgraph-checkpointer-couchbase
|
|
3
|
-
Version: 1.0.
|
|
3
|
+
Version: 1.0.7
|
|
4
4
|
Project-URL: Documentation, https://github.com/couchbase-ecosystem/langgraph-checkpointer-couchbase#readme
|
|
5
5
|
Project-URL: Issues, https://github.com/couchbase-ecosystem/langgraph-checkpointer-couchbase/issues
|
|
6
6
|
Project-URL: Source, https://github.com/couchbase-ecosystem/langgraph-checkpointer-couchbase
|
|
@@ -95,58 +95,72 @@ with CouchbaseSaver.from_conn_info(
|
|
|
95
95
|
bucket_name=os.getenv("CB_BUCKET") or "test",
|
|
96
96
|
scope_name=os.getenv("CB_SCOPE") or "langgraph",
|
|
97
97
|
) as checkpointer:
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
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
115
|
```
|
|
116
116
|
|
|
117
117
|
### Asynchronous Usage
|
|
118
118
|
|
|
119
119
|
```python
|
|
120
120
|
import os
|
|
121
|
+
from acouchbase.cluster import Cluster as ACluster
|
|
122
|
+
from couchbase.auth import PasswordAuthenticator
|
|
123
|
+
from couchbase.options import ClusterOptions
|
|
121
124
|
from langgraph_checkpointer_couchbase import AsyncCouchbaseSaver
|
|
122
125
|
from langgraph.graph import create_react_agent
|
|
123
126
|
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
127
|
+
auth = PasswordAuthenticator(
|
|
128
|
+
os.getenv("CB_USERNAME") or "Administrator",
|
|
129
|
+
os.getenv("CB_PASSWORD") or "password",
|
|
130
|
+
)
|
|
131
|
+
options = ClusterOptions(auth)
|
|
132
|
+
cluster = await ACluster.connect(os.getenv("CB_CLUSTER") or "couchbase://localhost", options)
|
|
133
|
+
|
|
134
|
+
bucket_name = os.getenv("CB_BUCKET") or "test"
|
|
135
|
+
scope_name = os.getenv("CB_SCOPE") or "langgraph"
|
|
136
|
+
|
|
137
|
+
async with AsyncCouchbaseSaver.from_cluster(
|
|
138
|
+
cluster=cluster,
|
|
139
|
+
bucket_name=bucket_name,
|
|
140
|
+
scope_name=scope_name,
|
|
130
141
|
) as checkpointer:
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
142
|
+
# Create the agent with checkpointing
|
|
143
|
+
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
|
|
144
|
+
|
|
145
|
+
# Configure with a unique thread ID
|
|
146
|
+
config = {"configurable": {"thread_id": "2"}}
|
|
147
|
+
|
|
148
|
+
# Run the agent asynchronously
|
|
149
|
+
res = await graph.ainvoke(
|
|
150
|
+
{"messages": [("human", "what's the weather in nyc")]}, config
|
|
151
|
+
)
|
|
152
|
+
|
|
153
|
+
# Retrieve checkpoints asynchronously
|
|
154
|
+
latest_checkpoint = await checkpointer.aget(config)
|
|
155
|
+
latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
|
|
156
|
+
checkpoint_tuples = [c async for c in checkpointer.alist(config)]
|
|
157
|
+
|
|
158
|
+
print(latest_checkpoint)
|
|
159
|
+
print(latest_checkpoint_tuple)
|
|
160
|
+
print(checkpoint_tuples)
|
|
161
|
+
|
|
162
|
+
# Close the cluster when done
|
|
163
|
+
await cluster.close()
|
|
150
164
|
```
|
|
151
165
|
|
|
152
166
|
## Configuration Options
|
|
@@ -68,58 +68,72 @@ with CouchbaseSaver.from_conn_info(
|
|
|
68
68
|
bucket_name=os.getenv("CB_BUCKET") or "test",
|
|
69
69
|
scope_name=os.getenv("CB_SCOPE") or "langgraph",
|
|
70
70
|
) as checkpointer:
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
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
88
|
```
|
|
89
89
|
|
|
90
90
|
### Asynchronous Usage
|
|
91
91
|
|
|
92
92
|
```python
|
|
93
93
|
import os
|
|
94
|
+
from acouchbase.cluster import Cluster as ACluster
|
|
95
|
+
from couchbase.auth import PasswordAuthenticator
|
|
96
|
+
from couchbase.options import ClusterOptions
|
|
94
97
|
from langgraph_checkpointer_couchbase import AsyncCouchbaseSaver
|
|
95
98
|
from langgraph.graph import create_react_agent
|
|
96
99
|
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
100
|
+
auth = PasswordAuthenticator(
|
|
101
|
+
os.getenv("CB_USERNAME") or "Administrator",
|
|
102
|
+
os.getenv("CB_PASSWORD") or "password",
|
|
103
|
+
)
|
|
104
|
+
options = ClusterOptions(auth)
|
|
105
|
+
cluster = await ACluster.connect(os.getenv("CB_CLUSTER") or "couchbase://localhost", options)
|
|
106
|
+
|
|
107
|
+
bucket_name = os.getenv("CB_BUCKET") or "test"
|
|
108
|
+
scope_name = os.getenv("CB_SCOPE") or "langgraph"
|
|
109
|
+
|
|
110
|
+
async with AsyncCouchbaseSaver.from_cluster(
|
|
111
|
+
cluster=cluster,
|
|
112
|
+
bucket_name=bucket_name,
|
|
113
|
+
scope_name=scope_name,
|
|
103
114
|
) as checkpointer:
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
115
|
+
# Create the agent with checkpointing
|
|
116
|
+
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
|
|
117
|
+
|
|
118
|
+
# Configure with a unique thread ID
|
|
119
|
+
config = {"configurable": {"thread_id": "2"}}
|
|
120
|
+
|
|
121
|
+
# Run the agent asynchronously
|
|
122
|
+
res = await graph.ainvoke(
|
|
123
|
+
{"messages": [("human", "what's the weather in nyc")]}, config
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
# Retrieve checkpoints asynchronously
|
|
127
|
+
latest_checkpoint = await checkpointer.aget(config)
|
|
128
|
+
latest_checkpoint_tuple = await checkpointer.aget_tuple(config)
|
|
129
|
+
checkpoint_tuples = [c async for c in checkpointer.alist(config)]
|
|
130
|
+
|
|
131
|
+
print(latest_checkpoint)
|
|
132
|
+
print(latest_checkpoint_tuple)
|
|
133
|
+
print(checkpoint_tuples)
|
|
134
|
+
|
|
135
|
+
# Close the cluster when done
|
|
136
|
+
await cluster.close()
|
|
123
137
|
```
|
|
124
138
|
|
|
125
139
|
## Configuration Options
|
|
@@ -1,12 +1,14 @@
|
|
|
1
1
|
from contextlib import asynccontextmanager
|
|
2
2
|
from datetime import timedelta
|
|
3
3
|
from typing import Any, AsyncIterator, Dict, Optional, Sequence, Tuple
|
|
4
|
+
import logging
|
|
4
5
|
|
|
5
6
|
from langchain_core.runnables import RunnableConfig
|
|
6
7
|
from acouchbase.cluster import Cluster as ACluster
|
|
7
8
|
from acouchbase.bucket import Bucket as ABucket
|
|
8
9
|
from couchbase.auth import PasswordAuthenticator
|
|
9
10
|
from couchbase.options import ClusterOptions, QueryOptions, UpsertOptions
|
|
11
|
+
from couchbase.exceptions import CollectionAlreadyExistsException
|
|
10
12
|
|
|
11
13
|
from langgraph.checkpoint.base import (
|
|
12
14
|
BaseCheckpointSaver,
|
|
@@ -18,6 +20,8 @@ from langgraph.checkpoint.base import (
|
|
|
18
20
|
)
|
|
19
21
|
from .utils import _encode_binary, _decode_binary
|
|
20
22
|
|
|
23
|
+
logger = logging.getLogger(__name__)
|
|
24
|
+
|
|
21
25
|
class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
22
26
|
"""A checkpoint saver that stores checkpoints in a Couchbase database."""
|
|
23
27
|
|
|
@@ -35,9 +39,35 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
35
39
|
self.cluster = cluster
|
|
36
40
|
self.bucket_name = bucket_name
|
|
37
41
|
self.scope_name = scope_name
|
|
42
|
+
self.bucket = self.cluster.bucket(bucket_name)
|
|
43
|
+
self.scope = self.bucket.scope(scope_name)
|
|
38
44
|
self.checkpoints_collection_name = checkpoints_collection_name
|
|
39
45
|
self.checkpoint_writes_collection_name = checkpoint_writes_collection_name
|
|
40
46
|
|
|
47
|
+
async def create_collections(self):
|
|
48
|
+
"""Create collections in the Couchbase bucket if they do not exist."""
|
|
49
|
+
|
|
50
|
+
collection_manager = self.bucket.collections()
|
|
51
|
+
try:
|
|
52
|
+
await collection_manager.create_collection(self.scope_name, self.checkpoints_collection_name)
|
|
53
|
+
except CollectionAlreadyExistsException as _:
|
|
54
|
+
pass
|
|
55
|
+
except Exception as e:
|
|
56
|
+
logger.exception("Error creating collections")
|
|
57
|
+
raise e
|
|
58
|
+
finally:
|
|
59
|
+
self.checkpoints_collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
|
|
60
|
+
|
|
61
|
+
try:
|
|
62
|
+
await collection_manager.create_collection(self.scope_name, self.checkpoint_writes_collection_name)
|
|
63
|
+
except CollectionAlreadyExistsException as _:
|
|
64
|
+
pass
|
|
65
|
+
except Exception as e:
|
|
66
|
+
logger.exception("Error creating collections")
|
|
67
|
+
raise e
|
|
68
|
+
finally:
|
|
69
|
+
self.checkpoint_writes_collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
|
|
70
|
+
|
|
41
71
|
@classmethod
|
|
42
72
|
@asynccontextmanager
|
|
43
73
|
async def from_conn_info(
|
|
@@ -69,15 +99,25 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
69
99
|
cls.bucket_name = bucket_name
|
|
70
100
|
cls.scope_name = scope_name
|
|
71
101
|
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
102
|
+
bucket = cluster.bucket(bucket_name)
|
|
103
|
+
await bucket.on_connect()
|
|
104
|
+
|
|
105
|
+
saver = AsyncCouchbaseSaver(
|
|
106
|
+
cluster,
|
|
107
|
+
bucket_name,
|
|
108
|
+
scope_name,
|
|
109
|
+
checkpoints_collection_name,
|
|
110
|
+
checkpoint_writes_collection_name,
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
await saver.create_collections()
|
|
75
114
|
|
|
76
115
|
yield saver
|
|
77
116
|
finally:
|
|
78
117
|
if cluster:
|
|
79
118
|
await cluster.close()
|
|
80
119
|
|
|
120
|
+
|
|
81
121
|
@classmethod
|
|
82
122
|
@asynccontextmanager
|
|
83
123
|
async def from_cluster(
|
|
@@ -98,9 +138,18 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
98
138
|
AsyncCouchbaseSaver: An instance of the AsyncCouchbaseSaver
|
|
99
139
|
"""
|
|
100
140
|
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
141
|
+
bucket = cluster.bucket(bucket_name)
|
|
142
|
+
await bucket.on_connect()
|
|
143
|
+
|
|
144
|
+
saver = AsyncCouchbaseSaver(
|
|
145
|
+
cluster,
|
|
146
|
+
bucket_name,
|
|
147
|
+
scope_name,
|
|
148
|
+
checkpoints_collection_name,
|
|
149
|
+
checkpoint_writes_collection_name,
|
|
150
|
+
)
|
|
151
|
+
|
|
152
|
+
await saver.create_collections()
|
|
104
153
|
|
|
105
154
|
yield saver
|
|
106
155
|
|
|
@@ -123,10 +172,10 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
123
172
|
checkpoint_id = get_checkpoint_id(config)
|
|
124
173
|
|
|
125
174
|
if checkpoint_id:
|
|
126
|
-
query = f'SELECT * FROM {self.bucket_name}
|
|
175
|
+
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
176
|
query_params = [thread_id, checkpoint_ns, checkpoint_id]
|
|
128
177
|
else:
|
|
129
|
-
query = f'SELECT * FROM {self.bucket_name}
|
|
178
|
+
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
179
|
query_params = [thread_id, checkpoint_ns]
|
|
131
180
|
|
|
132
181
|
result = self.cluster.query(query, QueryOptions(positional_parameters=query_params))
|
|
@@ -140,7 +189,7 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
140
189
|
}
|
|
141
190
|
checkpoint = self.serde.loads_typed((doc["type"], _decode_binary(doc["checkpoint"])))
|
|
142
191
|
|
|
143
|
-
serialized_writes_query = f'SELECT * FROM {self.bucket_name}
|
|
192
|
+
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'
|
|
144
193
|
serialized_writes_params = [thread_id, checkpoint_ns, doc["checkpoint_id"] or ""]
|
|
145
194
|
|
|
146
195
|
serialized_writes_result = self.cluster.query(serialized_writes_query, QueryOptions(positional_parameters=serialized_writes_params))
|
|
@@ -149,7 +198,7 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
149
198
|
async for write_doc in serialized_writes_result:
|
|
150
199
|
checkpoint_writes = write_doc.get(self.checkpoint_writes_collection_name, {})
|
|
151
200
|
if "task_id" not in checkpoint_writes:
|
|
152
|
-
|
|
201
|
+
logger.warning("'task_id' is not present in checkpoint_writes")
|
|
153
202
|
else:
|
|
154
203
|
pending_writes.append(
|
|
155
204
|
(
|
|
@@ -200,7 +249,7 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
200
249
|
AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
|
|
201
250
|
"""
|
|
202
251
|
|
|
203
|
-
query = f"SELECT * FROM {self.bucket_name}
|
|
252
|
+
query = f"SELECT * FROM `{self.bucket_name}`.`{self.scope_name}`.`{self.checkpoints_collection_name}` WHERE 1=1"
|
|
204
253
|
query_params = []
|
|
205
254
|
|
|
206
255
|
if config is not None:
|
|
@@ -294,6 +343,8 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
294
343
|
|
|
295
344
|
upsert_key = f"{thread_id}::{checkpoint_ns}::{checkpoint_id}"
|
|
296
345
|
|
|
346
|
+
# ensure bucket connected (idempotent)
|
|
347
|
+
await self.bucket.on_connect()
|
|
297
348
|
collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
|
|
298
349
|
await collection.upsert(upsert_key, (doc), UpsertOptions(timeout=timedelta(seconds=5)))
|
|
299
350
|
|
|
@@ -324,6 +375,7 @@ class AsyncCouchbaseSaver(BaseCheckpointSaver):
|
|
|
324
375
|
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
|
325
376
|
checkpoint_id = config["configurable"]["checkpoint_id"]
|
|
326
377
|
|
|
378
|
+
await self.bucket.on_connect()
|
|
327
379
|
collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
|
|
328
380
|
|
|
329
381
|
for idx, (channel, value) in enumerate(writes):
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
from contextlib import contextmanager
|
|
2
2
|
from datetime import timedelta
|
|
3
3
|
from typing import Any, Dict, Iterator, Optional, Sequence, Tuple
|
|
4
|
+
import logging
|
|
4
5
|
|
|
5
6
|
from langchain_core.runnables import RunnableConfig
|
|
6
7
|
from couchbase.cluster import Cluster
|
|
@@ -19,6 +20,8 @@ from langgraph.checkpoint.base import (
|
|
|
19
20
|
)
|
|
20
21
|
|
|
21
22
|
from .utils import _encode_binary, _decode_binary
|
|
23
|
+
|
|
24
|
+
logger = logging.getLogger(__name__)
|
|
22
25
|
class CouchbaseSaver(BaseCheckpointSaver):
|
|
23
26
|
"""A checkpoint saver that stores checkpoints in a Couchbase database.
|
|
24
27
|
|
|
@@ -118,7 +121,7 @@ class CouchbaseSaver(BaseCheckpointSaver):
|
|
|
118
121
|
except CollectionAlreadyExistsException as _:
|
|
119
122
|
pass
|
|
120
123
|
except Exception as e:
|
|
121
|
-
|
|
124
|
+
logger.exception("Error creating collections")
|
|
122
125
|
raise e
|
|
123
126
|
finally:
|
|
124
127
|
self.checkpoints_collection = self.bucket.scope(self.scope_name).collection(self.checkpoints_collection_name)
|
|
@@ -128,7 +131,7 @@ class CouchbaseSaver(BaseCheckpointSaver):
|
|
|
128
131
|
except CollectionAlreadyExistsException as _:
|
|
129
132
|
pass
|
|
130
133
|
except Exception as e:
|
|
131
|
-
|
|
134
|
+
logger.exception("Error creating collections")
|
|
132
135
|
raise e
|
|
133
136
|
finally:
|
|
134
137
|
self.checkpoint_writes_collection = self.bucket.scope(self.scope_name).collection(self.checkpoint_writes_collection_name)
|
|
@@ -153,10 +156,10 @@ class CouchbaseSaver(BaseCheckpointSaver):
|
|
|
153
156
|
checkpoint_id = get_checkpoint_id(config)
|
|
154
157
|
|
|
155
158
|
if checkpoint_id:
|
|
156
|
-
query = f'SELECT * FROM {self.bucket_name}
|
|
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'
|
|
157
160
|
query_params = [thread_id, checkpoint_ns, checkpoint_id]
|
|
158
161
|
else:
|
|
159
|
-
query = f'SELECT * FROM {self.bucket_name}
|
|
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'
|
|
160
163
|
query_params = [thread_id, checkpoint_ns]
|
|
161
164
|
|
|
162
165
|
result = self.cluster.query(query, QueryOptions(positional_parameters=query_params))
|
|
@@ -173,7 +176,7 @@ class CouchbaseSaver(BaseCheckpointSaver):
|
|
|
173
176
|
checkpoint_data = _decode_binary(doc["checkpoint"])
|
|
174
177
|
checkpoint = self.serde.loads_typed((doc["type"], checkpoint_data))
|
|
175
178
|
|
|
176
|
-
serialized_writes_query = f'SELECT * FROM {self.bucket_name}
|
|
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'
|
|
177
180
|
serialized_writes_params = [thread_id, checkpoint_ns, doc["checkpoint_id"] or ""]
|
|
178
181
|
serialized_writes_result = self.cluster.query(serialized_writes_query, QueryOptions(positional_parameters=serialized_writes_params))
|
|
179
182
|
|
|
@@ -181,7 +184,7 @@ class CouchbaseSaver(BaseCheckpointSaver):
|
|
|
181
184
|
for write_doc in serialized_writes_result:
|
|
182
185
|
checkpoint_writes = write_doc.get(self.checkpoint_writes_collection_name, {})
|
|
183
186
|
if "task_id" not in checkpoint_writes:
|
|
184
|
-
|
|
187
|
+
logger.warning("'task_id' is not present in checkpoint_writes")
|
|
185
188
|
else:
|
|
186
189
|
# Decode and deserialize value data
|
|
187
190
|
value_data = _decode_binary(checkpoint_writes["value"])
|
|
@@ -238,7 +241,7 @@ class CouchbaseSaver(BaseCheckpointSaver):
|
|
|
238
241
|
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
|
|
239
242
|
"""
|
|
240
243
|
|
|
241
|
-
query = f"SELECT * FROM {self.bucket_name}
|
|
244
|
+
query = f"SELECT * FROM `{self.bucket_name}`.`{self.scope_name}`.`{self.checkpoints_collection_name}` WHERE 1=1"
|
|
242
245
|
query_params = []
|
|
243
246
|
|
|
244
247
|
if config is not None:
|
|
@@ -4,7 +4,7 @@ from langchain_core.runnables import ConfigurableField
|
|
|
4
4
|
from langchain_core.tools import tool
|
|
5
5
|
from langchain_openai import ChatOpenAI
|
|
6
6
|
from langgraph.prebuilt import create_react_agent
|
|
7
|
-
from
|
|
7
|
+
from langgraph_checkpointer_couchbase import CouchbaseSaver, AsyncCouchbaseSaver
|
|
8
8
|
from dotenv import load_dotenv
|
|
9
9
|
import os
|
|
10
10
|
load_dotenv()
|
{langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/.env.example
RENAMED
|
File without changes
|
{langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/.gitignore
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{langgraph_checkpointer_couchbase-1.0.5 → langgraph_checkpointer_couchbase-1.0.7}/tests/__init__.py
RENAMED
|
File without changes
|