WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 21:27:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor: drop CreateCheckpointResult, caller-driven snapshot decision

Browse Source 
create_checkpoint now returns Checkpoint directly and accepts a
precomputed channels_to_snapshot set. Callers compute it via the
renamed delta_channels_to_snapshot helper and reuse it for counter
resets. Removes the NamedTuple wrapper and the .checkpoint boilerplate
at all 8 main.py call sites.
This commit is contained in:
Sydney Runkle
2026-05-07 12:05:53 -04:00
parent 506fc7eaf3
commit 9169af196c
3 changed files with 39 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/_checkpoint.py
+20 -36
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
from collections.abc import Callable, Mapping
from datetime import datetime, timezone
from typing import Any, NamedTuple, cast
from typing import Any, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
@@ -23,11 +23,6 @@ LATEST_VERSION = 4
GetNextVersion = Callable[[Any, None], Any]
class CreateCheckpointResult(NamedTuple):
checkpoint: Checkpoint
snapshotted: set[str]
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
v=LATEST_VERSION,
@@ -39,7 +34,7 @@ def empty_checkpoint() -> Checkpoint:
)
def decide_delta_snapshots(
def delta_channels_to_snapshot(
channels: Mapping[str, BaseChannel],
counts: Mapping[str, int],
) -> set[str]:
@@ -66,35 +61,30 @@ def create_checkpoint(
id: str | None = None,
updated_channels: set[str] | None = None,
get_next_version: GetNextVersion | None = None,
updates_since_snapshot: Mapping[str, int] | None = None,
) -> CreateCheckpointResult:
channels_to_snapshot: set[str] | None = None,
) -> Checkpoint:
"""Build a new Checkpoint from the previous one and live channel state.
For each `DeltaChannel`, a `_DeltaSnapshot(value)` blob is written into
`channel_values[k]` when `decide_delta_snapshots` says the channel should
snapshot (i.e. update count >= `snapshot_frequency`). Otherwise the
channel is omitted from `channel_values` and the ancestor walk
reconstructs state from `checkpoint_writes`.
Returns a `CreateCheckpointResult` containing the checkpoint and the
set of DeltaChannel names that were snapshotted (caller should reset
their per-channel counters to 0 for these).
For each name in `channels_to_snapshot`, a `_DeltaSnapshot(value)` blob
is written into `channel_values[k]`. Other delta channels are omitted
from `channel_values` — the ancestor walk reconstructs their state
from `checkpoint_writes`. Callers compute the set via
`delta_channels_to_snapshot(channels, counts)`; defaults to empty
(no snapshots) when not provided.
"""
ts = datetime.now(timezone.utc).isoformat()
counts = updates_since_snapshot or {}
snapshotted: set[str] = set()
channels_to_snapshot = channels_to_snapshot or set()
if channels is None:
values = checkpoint["channel_values"]
channel_versions = checkpoint["channel_versions"]
else:
will_snapshot = decide_delta_snapshots(channels, counts)
values = {}
channel_versions = dict(checkpoint["channel_versions"])
for k in channels:
if k not in channel_versions:
continue
ch = channels[k]
if k in will_snapshot:
if k in channels_to_snapshot:
# In exit mode, the snapshot decision is deferred to exit
# time (intermediate steps have do_checkpoint=False). The
# channel's count may have reached snapshot_frequency over
@@ -111,24 +101,18 @@ def create_checkpoint(
):
channel_versions[k] = get_next_version(channel_versions[k], None)
values[k] = _DeltaSnapshot(ch.get())
snapshotted.add(k)
else:
v = ch.checkpoint()
if v is not MISSING:
values[k] = v
return CreateCheckpointResult(
checkpoint=Checkpoint(
v=LATEST_VERSION,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=channel_versions,
versions_seen=checkpoint["versions_seen"],
updated_channels=None
if updated_channels is None
else sorted(updated_channels),
),
snapshotted=snapshotted,
return Checkpoint(
v=LATEST_VERSION,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=channel_versions,
versions_seen=checkpoint["versions_seen"],
updated_channels=None if updated_channels is None else sorted(updated_channels),
)
libs/langgraph/langgraph/pregel/_loop.py
+11 -7
View File
@@ -100,7 +100,7 @@ from langgraph.pregel._checkpoint import (
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
decide_delta_snapshots,
delta_channels_to_snapshot,
empty_checkpoint,
)
from langgraph.pregel._executor import (
@@ -1013,7 +1013,12 @@ class PregelLoop:
exiting or self.durability != "exit"
)
# create new checkpoint
result = create_checkpoint(
channels_to_snapshot = (
delta_channels_to_snapshot(self.channels, new_counts)
if do_checkpoint
else set()
)
self.checkpoint = create_checkpoint(
self.checkpoint,
self.channels if do_checkpoint else None,
self.step,
@@ -1022,10 +1027,9 @@ class PregelLoop:
get_next_version=self.checkpointer_get_next_version
if do_checkpoint
else None,
updates_since_snapshot=new_counts,
channels_to_snapshot=channels_to_snapshot,
)
self.checkpoint = result.checkpoint
for k in result.snapshotted:
for k in channels_to_snapshot:
new_counts[k] = 0
if new_counts:
self.checkpoint_metadata["delta_updates_since_snapshot"] = new_counts
@@ -1106,12 +1110,12 @@ class PregelLoop:
return
counts = self.checkpoint_metadata.get("delta_updates_since_snapshot", {}) or {}
will_snapshot = decide_delta_snapshots(self.channels, counts)
channels_to_snapshot = delta_channels_to_snapshot(self.channels, counts)
pending = [
(step, tid, ch, v)
for (step, tid, ch, v) in self._exit_delta_writes
if ch not in will_snapshot
if ch not in channels_to_snapshot
]
if not pending:
return
libs/langgraph/langgraph/pregel/main.py
+8 -8
View File
@@ -1726,7 +1726,7 @@ class Pregel(
# save checkpoint
next_config = checkpointer.put(
checkpoint_config,
create_checkpoint(checkpoint, channels, step).checkpoint,
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": step + 1,
@@ -1765,7 +1765,7 @@ class Pregel(
)
next_config = checkpointer.put(
checkpoint_config,
create_checkpoint(checkpoint, channels, next_step).checkpoint,
create_checkpoint(checkpoint, channels, next_step),
{
"source": "input",
"step": next_step,
@@ -1804,7 +1804,7 @@ class Pregel(
if saved is None:
raise InvalidUpdateError("Cannot copy a non-existent checkpoint")
next_checkpoint = create_checkpoint(checkpoint, None, step).checkpoint
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
next_config = checkpointer.put(
@@ -2009,7 +2009,7 @@ class Pregel(
checkpointer.get_next_version,
self.trigger_to_nodes,
)
checkpoint = create_checkpoint(checkpoint, channels, step + 1).checkpoint
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
next_config = checkpointer.put(
checkpoint_config,
checkpoint,
@@ -2175,7 +2175,7 @@ class Pregel(
# save checkpoint
next_config = await checkpointer.aput(
checkpoint_config,
create_checkpoint(checkpoint, channels, step).checkpoint,
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": step + 1,
@@ -2213,7 +2213,7 @@ class Pregel(
)
next_config = await checkpointer.aput(
checkpoint_config,
create_checkpoint(checkpoint, channels, next_step).checkpoint,
create_checkpoint(checkpoint, channels, next_step),
{
"source": "input",
"step": next_step,
@@ -2252,7 +2252,7 @@ class Pregel(
if saved is None:
raise InvalidUpdateError("Cannot copy a non-existent checkpoint")
next_checkpoint = create_checkpoint(checkpoint, None, step).checkpoint
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
next_config = await checkpointer.aput(
@@ -2456,7 +2456,7 @@ class Pregel(
checkpointer.get_next_version,
self.trigger_to_nodes,
)
checkpoint = create_checkpoint(checkpoint, channels, step + 1).checkpoint
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# save checkpoint, after applying writes
next_config = await checkpointer.aput(
checkpoint_config,