mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-12 12:47:53 +02:00
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user