WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-31 12:19:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

pass pending writes in checkpointers

Browse Source
This commit is contained in:
vbarda
2024-08-21 15:38:30 -04:00
parent a94168af7f
commit f51e7ea9a4
5 changed files with 33 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+1
View File
@@ -150,6 +150,7 @@ class PostgresSaver(BasePostgresSaver):
}
if value["parent_checkpoint_id"]
else None,
self._load_writes(value["pending_writes"]),
)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+1
View File
@@ -135,6 +135,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
}
if value["parent_checkpoint_id"]
else None,
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+15 -1
View File
@@ -318,7 +318,9 @@ class SqliteSaver(BaseCheckpointSaver):
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
with self.cursor(transaction=False) as cur:
with self.cursor(transaction=False) as cur, self.cursor(
transaction=False
) as writes_cur:
cur.execute(query, param_values)
for (
thread_id,
@@ -329,6 +331,14 @@ class SqliteSaver(BaseCheckpointSaver):
checkpoint,
metadata,
) in cur:
writes_cur.execute(
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
thread_id,
checkpoint_ns,
checkpoint_id,
),
)
yield CheckpointTuple(
{
"configurable": {
@@ -350,6 +360,10 @@ class SqliteSaver(BaseCheckpointSaver):
if parent_checkpoint_id
else None
),
[
(task_id, channel, self.serde.loads_typed((type, value)))
for task_id, channel, type, value in writes_cur
],
)
def put(
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+12
View File
@@ -346,6 +346,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
checkpoint,
metadata,
) in cursor:
writes_cur = await self.conn.execute(
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(
thread_id,
checkpoint_ns,
checkpoint_id,
),
)
yield CheckpointTuple(
{
"configurable": {
@@ -367,6 +375,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
if parent_checkpoint_id
else None
),
[
(task_id, channel, self.serde.loads_typed((type, value)))
async for task_id, channel, type, value in writes_cur
],
)
async def aput(
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+4
View File
@@ -206,6 +206,7 @@ class MemorySaver(
elif limit is not None:
limit -= 1
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
yield CheckpointTuple(
config={
"configurable": {
@@ -215,6 +216,9 @@ class MemorySaver(
}
},
checkpoint=self.serde.loads_typed(checkpoint),
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
metadata=metadata,
parent_config={
"configurable": {