Remove Checkpoint.pending_sends

- Instead store sends in a Topic channel, removing the need to fetch sends as writes against the parent checkpoint
- Remove deprecated/unused functions in langgraph-checkpoint (will require bumping min range for langgraph-checkpoint in langgraph lib)
- Implement migration of old pending sends in langgraph-checkpoint-postgres
- Ensure parent config of `checkpoint_during=False` checkpoints always points to checkpoints that were also saved
This commit is contained in:
Nuno Campos
2025-05-25 19:06:02 -07:00
parent 85522ec6a6
commit 4e8fbe4525
31 changed files with 559 additions and 508 deletions
@@ -1,4 +1,5 @@
import asyncio
from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import Any, Optional
@@ -130,7 +131,35 @@ class AsyncPostgresSaver(BasePostgresSaver):
# if we change this to use .stream() we need to make sure to close the cursor
async with self._cursor() as cur:
await cur.execute(query, args, binary=True)
async for value in cur:
values = await cur.fetchall()
if not values:
return
# migrate pending sends if necessary
if to_migrate := [
v
for v in values
if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
]:
await cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(
values[0]["thread_id"],
[v["parent_checkpoint_id"] for v in to_migrate],
),
)
grouped_by_parent = defaultdict(list)
for value in to_migrate:
grouped_by_parent[value["parent_checkpoint_id"]].append(value)
async for sends in cur:
for value in grouped_by_parent[sends["checkpoint_id"]]:
if value["channel_values"] is None:
value["channel_values"] = []
self._migrate_pending_sends(
sends["sends"],
value["checkpoint"],
value["channel_values"],
)
for value in values:
yield CheckpointTuple(
{
"configurable": {
@@ -143,7 +172,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
@@ -190,36 +218,52 @@ class AsyncPostgresSaver(BasePostgresSaver):
args,
binary=True,
)
value = await cur.fetchone()
if value is None:
return None
async for value in cur:
return CheckpointTuple(
# migrate pending sends if necessary
if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
await cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(thread_id, [value["parent_checkpoint_id"]]),
)
if sends := await cur.fetchone():
if value["channel_values"] is None:
value["channel_values"] = []
self._migrate_pending_sends(
sends["sends"],
value["checkpoint"],
value["channel_values"],
)
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
"checkpoint_id": value["parent_checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None
),
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
}
if value["parent_checkpoint_id"]
else None
),
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
async def aput(
self,