WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-30 03:39:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove Checkpoint.pending_sends

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/README.md
+2 -4
View File
@@ -25,7 +25,7 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
# call .setup() the first time you're using the checkpointer
checkpointer.setup()
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -47,7 +47,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
@@ -67,7 +66,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -89,7 +88,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+70 -27
View File
@@ -1,4 +1,5 @@
import threading
from collections import defaultdict
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import Any, Optional
@@ -142,8 +143,36 @@ class PostgresSaver(BasePostgresSaver):
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
with self._cursor() as cur:
cur.execute(query, args, binary=True)
for value in cur:
cur.execute(query, args)
values = 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"]
]:
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)
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": {
@@ -155,7 +184,6 @@ class PostgresSaver(BasePostgresSaver):
self._load_checkpoint(
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
@@ -221,37 +249,52 @@ class PostgresSaver(BasePostgresSaver):
cur.execute(
self.SELECT_SQL + where,
args,
binary=True,
)
value = cur.fetchone()
if value is None:
return None
for value in cur:
return CheckpointTuple(
# migrate pending sends if necessary
if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(thread_id, [value["parent_checkpoint_id"]]),
)
if sends := 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"],
}
},
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"],
}
},
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
),
self._load_writes(value["pending_writes"]),
)
}
if value["parent_checkpoint_id"]
else None
),
self._load_writes(value["pending_writes"]),
)
def put(
self,
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+70 -26
View File
@@ -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,
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+37 -17
View File
@@ -14,7 +14,7 @@ from langgraph.checkpoint.base import (
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
from langgraph.checkpoint.serde.types import TASKS
MetadataInput = Optional[dict[str, Any]]
@@ -72,7 +72,7 @@ MIGRATIONS = [
"""ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';""",
]
SELECT_SQL = f"""
SELECT_SQL = """
select
thread_id,
checkpoint,
@@ -96,17 +96,20 @@ select
where cw.thread_id = checkpoints.thread_id
and cw.checkpoint_ns = checkpoints.checkpoint_ns
and cw.checkpoint_id = checkpoints.checkpoint_id
) as pending_writes,
(
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx)
from checkpoint_writes cw
where cw.thread_id = checkpoints.thread_id
and cw.checkpoint_ns = checkpoints.checkpoint_ns
and cw.checkpoint_id = checkpoints.parent_checkpoint_id
and cw.channel = '{TASKS}'
) as pending_sends
) as pending_writes
from checkpoints """
SELECT_PENDING_SENDS_SQL = f"""
select
checkpoint_id,
array_agg(array[type::bytea, blob] order by task_path, task_id, idx) as sends
from checkpoint_writes
where thread_id = %s
and checkpoint_id = any(%s)
and channel = '{TASKS}'
group by checkpoint_id
"""
UPSERT_CHECKPOINT_BLOBS_SQL = """
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
VALUES (%s, %s, %s, %s, %s, %s)
@@ -140,6 +143,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
MIGRATIONS = MIGRATIONS
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
@@ -149,22 +153,38 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
jsonplus_serde = JsonPlusSerializer()
supports_pipeline: bool
def _migrate_pending_sends(
self,
pending_sends: list[tuple[bytes, bytes]],
checkpoint: dict[str, Any],
channel_values: list[tuple[bytes, bytes, bytes]],
) -> None:
if not pending_sends:
return
# add to values
enc, blob = self.serde.dumps_typed(
[self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends],
)
channel_values.append((TASKS.encode(), enc.encode(), blob))
# add to versions
checkpoint["channel_versions"][TASKS] = (
max(checkpoint["channel_versions"].values())
if checkpoint["channel_versions"]
else self.get_next_version(None)
)
def _load_checkpoint(
self,
checkpoint: dict[str, Any],
channel_values: list[tuple[bytes, bytes, bytes]],
pending_sends: list[tuple[bytes, bytes]],
) -> Checkpoint:
return {
**checkpoint,
"pending_sends": [
self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
],
"channel_values": self._load_blobs(channel_values),
}
def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
return {**checkpoint, "pending_sends": []}
return checkpoint
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
@@ -249,7 +269,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
# NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
return serialized_metadata.decode().replace("\\u0000", "")
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
def get_next_version(self, current: Optional[str]) -> str:
if current is None:
current_v = 0
elif isinstance(current, int):
libs/checkpoint-postgres/tests/checkpoint_utils.py
+51
View File
@@ -0,0 +1,51 @@
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Any, Optional, Protocol
from langgraph.checkpoint.base import Checkpoint, EmptyChannelError
from langgraph.checkpoint.base.id import uuid6
class ChannelProtocol(Protocol):
def checkpoint(self) -> Optional[Any]: ...
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
v=1,
id=str(uuid6(clock_seq=-2)),
ts=datetime.now(timezone.utc).isoformat(),
channel_values={},
channel_versions={},
versions_seen={},
)
def create_checkpoint(
checkpoint: Checkpoint,
channels: Optional[Mapping[str, ChannelProtocol]],
step: int,
*,
id: Optional[str] = None,
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
if channels is None:
values = checkpoint["channel_values"]
else:
values = {}
for k, v in channels.items():
if k not in checkpoint["channel_versions"]:
continue
try:
values[k] = v.checkpoint()
except EmptyChannelError:
pass
return Checkpoint(
v=1,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
)
libs/checkpoint-postgres/tests/test_async.py
+51 -2
View File
@@ -14,10 +14,10 @@ from langgraph.checkpoint.base import (
EXCLUDED_METADATA_KEYS,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.serde.types import TASKS
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -263,3 +263,52 @@ async def test_null_chars(saver_name: str, test_data) -> None:
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
0
].metadata["my_key"] == "abc"
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_pending_sends_migration(saver_name: str) -> None:
async with _saver(saver_name) as saver:
config = {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
}
}
# create the first checkpoint
# and put some pending sends
checkpoint_0 = empty_checkpoint()
config = await saver.aput(config, checkpoint_0, {}, {})
await saver.aput_writes(
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
)
await saver.aput_writes(config, [(TASKS, "send-3")], task_id="task-2")
# check that fetching checkpoint_0 doesn't attach pending sends
# (they should be attached to the next checkpoint)
tuple_0 = await saver.aget_tuple(config)
assert tuple_0.checkpoint["channel_values"] == {}
assert tuple_0.checkpoint["channel_versions"] == {}
# create the second checkpoint
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
config = await saver.aput(config, checkpoint_1, {}, {})
# check that pending sends are attached to checkpoint_1
tuple_1 = await saver.aget_tuple(config)
assert tuple_1.checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in tuple_1.checkpoint["channel_versions"]
# check that list also applies the migration
search_results = [
c async for c in saver.alist({"configurable": {"thread_id": "thread-1"}})
]
assert len(search_results) == 2
assert search_results[-1].checkpoint["channel_values"] == {}
assert search_results[-1].checkpoint["channel_versions"] == {}
assert search_results[0].checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in search_results[0].checkpoint["channel_versions"]
libs/checkpoint-postgres/tests/test_sync.py
+51 -2
View File
@@ -15,10 +15,10 @@ from langgraph.checkpoint.base import (
EXCLUDED_METADATA_KEYS,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.serde.types import TASKS
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -258,3 +258,52 @@ def test_nonnull_migrations() -> None:
for migration in PostgresSaver.MIGRATIONS:
statement = _leading_comment_remover.sub("", migration).split()[0]
assert statement.strip()
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
def test_pending_sends_migration(saver_name: str) -> None:
with _saver(saver_name) as saver:
config = {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
}
}
# create the first checkpoint
# and put some pending sends
checkpoint_0 = empty_checkpoint()
config = saver.put(config, checkpoint_0, {}, {})
saver.put_writes(
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
)
saver.put_writes(config, [(TASKS, "send-3")], task_id="task-2")
# check that fetching checkpoint_0 doesn't attach pending sends
# (they should be attached to the next checkpoint)
tuple_0 = saver.get_tuple(config)
assert tuple_0.checkpoint["channel_values"] == {}
assert tuple_0.checkpoint["channel_versions"] == {}
# create the second checkpoint
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
config = saver.put(config, checkpoint_1, {}, {})
# check that pending sends are attached to checkpoint_1
checkpoint_1 = saver.get_tuple(config)
assert checkpoint_1.checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in checkpoint_1.checkpoint["channel_versions"]
# check that list also applies the migration
search_results = [
c for c in saver.list({"configurable": {"thread_id": "thread-1"}})
]
assert len(search_results) == 2
assert search_results[-1].checkpoint["channel_values"] == {}
assert search_results[-1].checkpoint["channel_versions"] == {}
assert search_results[0].checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in search_results[0].checkpoint["channel_versions"]