mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-05 09:17:47 +02:00
feat(channels): step-based eager snapshot_frequency via _DeltaSnapshot ext type
Replaces the write-count (_write_count) snapshot mechanism with a clean step-based approach: pregel's create_checkpoint fires snapshots every N pregel steps regardless of whether the channel was written (eager). Key changes: - snapshot_frequency=None (default) for pure delta; int N for snapshot every N steps - DeltaChannel.checkpoint() always returns DELTA_SENTINEL; snapshot logic lives in create_checkpoint which has the step number - create_checkpoint bumps channel version via get_next_version when the channel wasn't written at a snapshot step (eager: always stores the blob) - _DeltaSnapshot NamedTuple registered as EXT_DELTA_SNAPSHOT (code 7) in the msgpack serde — no dict key collision, type tag does the dispatch - InMemorySaver and PostgresSaver _get_channel_writes_history updated: _DeltaSnapshot blobs collect pending_writes before terminating (they encode the NEXT step's transition, not subsumed by the snapshot unlike pre-delta migration blobs) Tests confirm: - Snapshots fire at every N steps even when channel has no write that step - Correct accumulated state after reconstruction from snapshot + replay Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
af9ef2a1f9
commit
bf30da1c8f
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
MetadataInput = dict[str, Any] | None
|
||||
@@ -281,13 +281,24 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
collected: list[PendingWrite] = [] # newest first; reversed at the end
|
||||
for cid in ancestors:
|
||||
# Pre-delta blob terminator: subsumes any writes at this ancestor.
|
||||
ver = ver_of.get(cid)
|
||||
if ver is not None:
|
||||
seed_blob = blob_by_ver.get(ver)
|
||||
if seed_blob is not None and seed_blob[0] != "empty":
|
||||
blob_value = self.serde.loads_typed(seed_blob)
|
||||
if blob_value is not DELTA_SENTINEL:
|
||||
if isinstance(blob_value, _DeltaSnapshot):
|
||||
# Step-based snapshot: collect this ancestor's
|
||||
# pending_writes first (they encode the NEXT step's
|
||||
# transition, not subsumed by the snapshot blob).
|
||||
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(
|
||||
cid, []
|
||||
):
|
||||
val = self.serde.loads_typed((type_tag, write_blob))
|
||||
collected.append((task_id, channel, val))
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=blob_value, writes=collected)
|
||||
# Pre-delta blob: subsumes this ancestor's writes.
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=blob_value, writes=collected)
|
||||
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
|
||||
|
||||
@@ -27,6 +27,7 @@ from langgraph.checkpoint.base import (
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -185,8 +186,31 @@ class InMemorySaver(
|
||||
if blob_entry is not None and blob_entry[0] != "empty":
|
||||
blob_value = self.serde.loads_typed(blob_entry)
|
||||
if blob_value is not DELTA_SENTINEL:
|
||||
# Pre-delta snapshot terminator. Skip this
|
||||
# ancestor's writes — the blob subsumes them.
|
||||
if isinstance(blob_value, _DeltaSnapshot):
|
||||
# Step-based snapshot: the blob is state AT this
|
||||
# ancestor, but the ancestor's pending_writes
|
||||
# encode the NEXT step's transition and are NOT
|
||||
# subsumed by the snapshot — collect them first.
|
||||
step_writes = self.writes.get(
|
||||
(thread_id, checkpoint_ns, cp_id), {}
|
||||
)
|
||||
for (_task_id, _idx), (
|
||||
tid,
|
||||
ch,
|
||||
serialized,
|
||||
_,
|
||||
) in sorted(step_writes.items(), reverse=True):
|
||||
if ch != channel:
|
||||
continue
|
||||
collected.append(
|
||||
(tid, ch, self.serde.loads_typed(serialized))
|
||||
)
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
)
|
||||
# Pre-delta blob: state AT this ancestor already
|
||||
# subsumes its pending_writes — skip them.
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
|
||||
@@ -33,7 +33,7 @@ from langchain_core.load.load import Reviver
|
||||
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
|
||||
from langgraph.checkpoint.serde.types import DELTA_SENTINEL, SendProtocol
|
||||
from langgraph.checkpoint.serde.types import DELTA_SENTINEL, SendProtocol, _DeltaSnapshot
|
||||
from langgraph.store.base import Item
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -296,10 +296,13 @@ EXT_METHOD_SINGLE_ARG = 3
|
||||
EXT_PYDANTIC_V1 = 4
|
||||
EXT_PYDANTIC_V2 = 5
|
||||
EXT_NUMPY_ARRAY = 6
|
||||
EXT_DELTA_SNAPSHOT = 7
|
||||
|
||||
|
||||
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
if hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
||||
if isinstance(obj, _DeltaSnapshot):
|
||||
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
|
||||
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
||||
return ormsgpack.Ext(
|
||||
EXT_PYDANTIC_V2,
|
||||
_msgpack_enc(
|
||||
@@ -613,7 +616,13 @@ def _create_msgpack_ext_hook(
|
||||
return False
|
||||
|
||||
def ext_hook(code: int, data: bytes) -> Any:
|
||||
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
if code == EXT_DELTA_SNAPSHOT:
|
||||
return _DeltaSnapshot(
|
||||
ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
)
|
||||
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import (
|
||||
Any,
|
||||
NamedTuple,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
runtime_checkable,
|
||||
@@ -33,6 +34,20 @@ class _DeltaSentinel:
|
||||
DELTA_SENTINEL = _DeltaSentinel()
|
||||
|
||||
|
||||
class _DeltaSnapshot(NamedTuple):
|
||||
"""Snapshot blob for a DeltaChannel with finite snapshot_frequency.
|
||||
|
||||
Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code.
|
||||
The ancestor walk in `_get_channel_writes_history` terminates when it
|
||||
encounters this type (any non-sentinel blob stops the walk).
|
||||
|
||||
`from_checkpoint` reconstructs the channel value directly from `.value`
|
||||
without replaying writes — the snapshot IS the accumulated state.
|
||||
"""
|
||||
|
||||
value: Any
|
||||
|
||||
|
||||
Value = TypeVar("Value", covariant=True)
|
||||
Update = TypeVar("Update", contravariant=True)
|
||||
C = TypeVar("C")
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy as _copy
|
||||
import math
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._constants import OVERWRITE
|
||||
@@ -40,38 +40,34 @@ def _get_overwrite(value: Any) -> tuple[bool, Any]:
|
||||
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
"""Fold-reducer channel with configurable snapshot cadence.
|
||||
|
||||
`snapshot_frequency=math.inf` (default): pure delta mode — stores only
|
||||
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes
|
||||
through `operator`. O(N) storage, O(N) read replay depth.
|
||||
`snapshot_frequency=None` (default): pure delta — stores only
|
||||
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes.
|
||||
|
||||
`snapshot_frequency=N` (integer > 1): writes a full snapshot blob every
|
||||
N steps; on non-snapshot steps stores `DELTA_SENTINEL`. Reads walk at
|
||||
most N ancestors before finding the snapshot, bounding replay depth to N.
|
||||
Storage is O(N²/N) = O(N·avg_msg_size) — still linear per step,
|
||||
but snapshots grow with accumulated messages so total is O(N²/N).
|
||||
`snapshot_frequency=N`: pregel's `create_checkpoint` writes a full
|
||||
`_DeltaSnapshot` blob every N steps (eagerly, even if the channel had
|
||||
no write that step). Reads walk at most N ancestor checkpoints before
|
||||
hitting the snapshot, bounding replay depth to N regardless of thread
|
||||
length.
|
||||
|
||||
Parameters:
|
||||
operator: Binary reducer `(Value, Value) -> Value` applied pairwise.
|
||||
Must be associative when `snapshot_frequency > 1` since replayed
|
||||
writes fold through the same operator as live writes.
|
||||
snapshot_frequency: Every Nth step writes a full snapshot blob.
|
||||
Default `math.inf` (pure delta, never snapshot).
|
||||
operator: Binary reducer `(Value, Value) -> Value`.
|
||||
snapshot_frequency: Every Nth pregel step writes a snapshot blob.
|
||||
`None` (default) = pure delta, never snapshot.
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "operator", "snapshot_frequency", "_write_count")
|
||||
__slots__ = ("value", "operator", "snapshot_frequency")
|
||||
value: Value | Any
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Any, Any], Any],
|
||||
*,
|
||||
snapshot_frequency: int | float = math.inf,
|
||||
snapshot_frequency: int | None = None,
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.snapshot_frequency = snapshot_frequency
|
||||
self.value: Any = []
|
||||
self._write_count: int = 0
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
@@ -93,19 +89,23 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
def is_snapshot_step(self, step: int) -> bool:
|
||||
"""True if pregel should write a snapshot blob at this step."""
|
||||
return (
|
||||
self.snapshot_frequency is not None and step % self.snapshot_frequency == 0
|
||||
)
|
||||
|
||||
def _clone_empty(self) -> Self:
|
||||
new = self.__class__.__new__(self.__class__)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
new.operator = self.operator
|
||||
new.snapshot_frequency = self.snapshot_frequency
|
||||
new._write_count = 0
|
||||
new.value = MISSING
|
||||
return new
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = self._clone_empty()
|
||||
new._write_count = self._write_count
|
||||
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
||||
return new
|
||||
|
||||
@@ -123,36 +123,24 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
"""Initialize from a stored blob or sentinel.
|
||||
|
||||
Blob formats:
|
||||
* `DELTA_SENTINEL` / `MISSING`: start empty; caller will replay writes.
|
||||
* `{"__delta_v__": value, "__delta_wc__": n}`: snapshot blob — restore
|
||||
value and write count so future writes trigger the next snapshot at
|
||||
the right cadence.
|
||||
* plain value (migration from old DeltaChannel blobs): restore value,
|
||||
reset write count to 0.
|
||||
Blob types (dispatched via serde ext code, not dict key inspection):
|
||||
* `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes.
|
||||
* `_DeltaSnapshot(value)`: restore value directly from snapshot.
|
||||
* plain value (migration from old BinOp blobs): use directly.
|
||||
"""
|
||||
new = self._clone_empty()
|
||||
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
|
||||
new.value = _empty(self.typ)
|
||||
new._write_count = 0
|
||||
elif isinstance(checkpoint, dict) and "__delta_v__" in checkpoint:
|
||||
new.value = checkpoint["__delta_v__"]
|
||||
new._write_count = checkpoint["__delta_wc__"]
|
||||
new.value = _empty(new.typ)
|
||||
elif isinstance(checkpoint, _DeltaSnapshot):
|
||||
new.value = checkpoint.value
|
||||
else:
|
||||
new.value = checkpoint
|
||||
new._write_count = 0
|
||||
return new
|
||||
|
||||
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
|
||||
"""Fold ancestor writes oldest→newest into current value.
|
||||
|
||||
Also increments `_write_count` per replayed write so the snapshot
|
||||
cadence stays correct across invocations. The count resumes from
|
||||
wherever the seed's `__delta_wc__` left off (set by `from_checkpoint`).
|
||||
"""
|
||||
"""Fold ancestor writes oldest→newest into current value."""
|
||||
for _, _, value in writes:
|
||||
self.value = self._apply_write(self.value, value)
|
||||
self._write_count += 1
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
if not values:
|
||||
@@ -171,7 +159,6 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
elif seen_overwrite:
|
||||
continue
|
||||
self.value = self._apply_write(self.value, value)
|
||||
self._write_count += 1
|
||||
return True
|
||||
|
||||
def get(self) -> Any:
|
||||
@@ -183,19 +170,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
"""Return stored representation.
|
||||
"""Return stored representation: always `DELTA_SENTINEL`.
|
||||
|
||||
Pure delta mode (`snapshot_frequency=math.inf`) always returns
|
||||
`DELTA_SENTINEL`. For finite `snapshot_frequency`, returns a snapshot
|
||||
blob `{"__delta_v__": value, "__delta_wc__": n}` every
|
||||
`snapshot_frequency` writes, with the current write count embedded so
|
||||
`from_checkpoint` can restore the counter and maintain correct cadence
|
||||
across invocations. All other writes return `DELTA_SENTINEL`.
|
||||
Snapshot decisions are made by `create_checkpoint` in pregel (which
|
||||
has the step number) via `is_snapshot_step`. `checkpoint()` is only
|
||||
called for non-snapshot steps or when no checkpointer is available.
|
||||
"""
|
||||
if self.value is MISSING:
|
||||
return MISSING
|
||||
if self.snapshot_frequency == math.inf:
|
||||
return DELTA_SENTINEL
|
||||
if self._write_count > 0 and self._write_count % self.snapshot_frequency == 0:
|
||||
return {"__delta_v__": self.value, "__delta_wc__": self._write_count}
|
||||
return DELTA_SENTINEL
|
||||
|
||||
@@ -1,19 +1,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel # used in _needs_replay
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
LATEST_VERSION = 4
|
||||
|
||||
GetNextVersion = Callable[[Any, None], Any]
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
@@ -33,31 +37,50 @@ def create_checkpoint(
|
||||
*,
|
||||
id: str | None = None,
|
||||
updated_channels: set[str] | None = None,
|
||||
get_next_version: GetNextVersion | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels.
|
||||
|
||||
`DeltaChannel.checkpoint()` returns `DELTA_SENTINEL` or the full value
|
||||
depending on its internal write count, so no special handling is needed
|
||||
here. The snapshot is stored at the exact step of the triggering write.
|
||||
For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a
|
||||
`_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor
|
||||
walk to at most N steps. Snapshots are eager: even if the channel had no
|
||||
write this step, a version bump is forced (via `get_next_version`) so the
|
||||
blob is stored by `put()`. Without `get_next_version` (e.g. static
|
||||
contexts), snapshot steps gracefully fall back to sentinel.
|
||||
"""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
channel_versions = checkpoint["channel_versions"]
|
||||
else:
|
||||
values = {}
|
||||
channel_versions = dict(checkpoint["channel_versions"])
|
||||
for k in channels:
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
if k not in channel_versions:
|
||||
continue
|
||||
ch = channels[k]
|
||||
v = ch.checkpoint()
|
||||
if v is not MISSING:
|
||||
values[k] = v
|
||||
if (
|
||||
isinstance(ch, DeltaChannel)
|
||||
and ch.is_snapshot_step(step)
|
||||
and ch.is_available()
|
||||
):
|
||||
# Eager snapshot: bump version if not already written this step
|
||||
# so put() includes this channel in new_versions and stores blob.
|
||||
if get_next_version is not None and (
|
||||
updated_channels is None or k not in updated_channels
|
||||
):
|
||||
channel_versions[k] = get_next_version(channel_versions[k], None)
|
||||
values[k] = _DeltaSnapshot(ch.get())
|
||||
else:
|
||||
v = ch.checkpoint()
|
||||
if v is not MISSING:
|
||||
values[k] = v
|
||||
return Checkpoint(
|
||||
v=LATEST_VERSION,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
channel_versions=channel_versions,
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
updated_channels=None if updated_channels is None else sorted(updated_channels),
|
||||
)
|
||||
@@ -67,9 +90,8 @@ def _needs_replay(spec: BaseChannel, stored: object) -> bool:
|
||||
"""True if `spec` is a `DeltaChannel` and the stored blob is a sentinel,
|
||||
requiring an ancestor walk to reconstruct.
|
||||
|
||||
Snapshot blobs are dicts with `__delta_v__`; sentinels are `DELTA_SENTINEL`.
|
||||
Plain non-sentinel values are migration blobs (old DeltaChannel format).
|
||||
All non-sentinel values resolve directly via `from_checkpoint`.
|
||||
`_DeltaSnapshot` blobs and plain values (migration) resolve directly via
|
||||
`from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay.
|
||||
"""
|
||||
if not isinstance(spec, DeltaChannel):
|
||||
return False
|
||||
@@ -86,16 +108,10 @@ def channels_from_checkpoint(
|
||||
"""Hydrate channels from a checkpoint.
|
||||
|
||||
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
|
||||
is sufficient — the stored value IS the reconstructed state.
|
||||
|
||||
`DeltaChannel` is the exception: its stored value on non-snapshot steps is
|
||||
`DELTA_SENTINEL`; the full state is spread across `checkpoint_writes` along
|
||||
the ancestor chain. When `saver` and `config` are provided, this function
|
||||
fetches that history via `saver._get_channel_writes_history` and folds it
|
||||
through the channel's operator. The walk stops at the nearest full-snapshot
|
||||
blob, so read depth is bounded by `snapshot_frequency`. Without saver/config
|
||||
(static contexts — graph drawing, unit tests), delta-mode channels fall back
|
||||
to empty.
|
||||
is sufficient. `DeltaChannel` is the exception: sentinel blobs require an
|
||||
ancestor walk via `saver._get_channel_writes_history`. The walk terminates
|
||||
at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain
|
||||
value, so read depth is bounded by `snapshot_frequency`.
|
||||
"""
|
||||
channel_specs: dict[str, BaseChannel] = {}
|
||||
managed_specs: dict[str, ManagedValueSpec] = {}
|
||||
@@ -110,10 +126,6 @@ def channels_from_checkpoint(
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
if _needs_replay(spec, stored) and saver is not None and config is not None:
|
||||
# Walk ancestors for seed + writes. The saver's walk stops at
|
||||
# the nearest non-sentinel blob (natural terminator under
|
||||
# snapshot_frequency > 1; pre-migration blobs also act as
|
||||
# terminators if the spec was changed mid-thread).
|
||||
assert isinstance(spec, DeltaChannel)
|
||||
history = saver._get_channel_writes_history(config, k)
|
||||
replay_ch = spec.from_checkpoint(history.seed)
|
||||
|
||||
@@ -891,6 +891,9 @@ class PregelLoop:
|
||||
self.step,
|
||||
id=self.checkpoint["id"] if exiting else None,
|
||||
updated_channels=self.updated_channels,
|
||||
get_next_version=self.checkpointer_get_next_version
|
||||
if do_checkpoint
|
||||
else None,
|
||||
)
|
||||
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
|
||||
if TASKS in self.checkpoint["channel_values"] and any(
|
||||
|
||||
@@ -275,6 +275,124 @@ def test_delta_channel_checkpoint_returns_sentinel() -> None:
|
||||
assert ch.checkpoint() is DELTA_SENTINEL
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_step_based() -> None:
|
||||
"""Snapshots fire on every Nth step regardless of whether the channel was written.
|
||||
|
||||
With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot
|
||||
blob — even if the channel had no write that step (eager snapshot). This
|
||||
bounds the ancestor walk to at most N steps on any read.
|
||||
"""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# snapshot_frequency=5: snapshot every 5 pregel steps
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_frequency=5)]
|
||||
other: str
|
||||
|
||||
def node_a(state: State) -> dict:
|
||||
# writes to messages
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
||||
|
||||
def node_b(state: State) -> dict:
|
||||
# writes ONLY to other, not messages — snapshot must still fire at step N
|
||||
return {"other": "y"}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("a", node_a)
|
||||
g.add_node("b", node_b)
|
||||
g.add_edge(START, "a")
|
||||
g.add_edge("a", "b")
|
||||
saver = InMemorySaver()
|
||||
graph = g.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
for i in range(6):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""},
|
||||
config,
|
||||
)
|
||||
|
||||
# Confirm at least one snapshot blob exists for messages
|
||||
msg_blob_values = [
|
||||
saver.serde.loads_typed((type_tag, blob))
|
||||
for k, (type_tag, blob) in saver.blobs.items()
|
||||
if k[2] == "messages" and type_tag == "msgpack" and blob
|
||||
]
|
||||
snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)]
|
||||
assert snapshots, "expected at least one _DeltaSnapshot blob for messages"
|
||||
|
||||
# Final state must be correct regardless of snapshot cadence
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 12 # 6 human + 6 AI
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_fires_even_when_not_written() -> None:
|
||||
"""Eager snapshot: _DeltaSnapshot stored at snapshot step even when the
|
||||
channel had no write that step (node_b doesn't touch messages).
|
||||
"""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_frequency=3)]
|
||||
tick: int
|
||||
|
||||
def writer(state: State) -> dict:
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
||||
|
||||
def ticker(state: State) -> dict:
|
||||
# never writes messages
|
||||
return {"tick": state["tick"] + 1}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("writer", writer)
|
||||
g.add_node("ticker", ticker)
|
||||
g.add_edge(START, "writer")
|
||||
g.add_edge("writer", "ticker")
|
||||
saver = InMemorySaver()
|
||||
graph = g.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
for i in range(5):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0},
|
||||
config,
|
||||
)
|
||||
|
||||
# Count distinct message channel blob versions
|
||||
msg_blobs = {
|
||||
k: saver.serde.loads_typed((t, b))
|
||||
for k, (t, b) in saver.blobs.items()
|
||||
if k[2] == "messages" and t == "msgpack" and b
|
||||
}
|
||||
snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)}
|
||||
# There must be snapshots (ticker steps are snapshot steps too)
|
||||
assert snapshots, (
|
||||
"eager snapshots must fire even on steps where messages wasn't written"
|
||||
)
|
||||
|
||||
# All get_state calls must return the correct accumulated value
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 10 # 5 human + 5 AI
|
||||
|
||||
|
||||
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
|
||||
from typing import Annotated
|
||||
|
||||
Reference in New Issue
Block a user