mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 02:37:52 +02:00
refactor(delta-channel): drop snapshot_every and saver Overwrite terminator
snapshot_every was a knob for bounding reconstruction cost on deep threads. Benchmarks (notes/add_messages_replay_problem.md + scratch work on sr/add-messages-replay-bench) showed the add_messages fast-path (optimize/add-messages-fast-path) closes the quadratic replay cost for threads under ~1000 turns, where the crossover to snapshots makes sense. For deeper threads we'll ship a first-class compaction primitive instead. Removals: * DeltaChannel: snapshot_every ctor param, _writes_since_snapshot counter, should_snapshot() / snapshot_write() methods, counter threading through _apply_write / update / from_checkpoint / copy. * Pregel loop: post-checkpoint snapshot-injection block and SNAPSHOT_TASK_ID import + constant. * Checkpoint base: _overwrite_types() helper and the ancestor-walk short-circuit on user-emitted Overwrite in sync + async get_channel_writes. * InMemory + Postgres savers: same walk-terminator shortcut. The pre-delta blob terminator (seed-from-ancestor-blob) stays — it's required for migration correctness, not a snapshot optimization. * Tests for all of the above. Preserved: * Channel-level Overwrite semantics in DeltaChannel / BinOpAggregate: Overwrite still resets the value at reducer level; same-super-step dedup and InvalidUpdateError on multiple Overwrites still enforced. * Pre-delta migration seeding.
This commit is contained in:
Sydney Runkle
@@ -13,7 +13,6 @@ from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
DeltaChannelWrites,
|
||||
_overwrite_types,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
@@ -238,19 +237,14 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
(`AsyncPostgresSaver`); both paths run the queries themselves and
|
||||
feed the rows here.
|
||||
|
||||
Walk is newest → oldest from the target's parent. Stops at the first
|
||||
terminator:
|
||||
|
||||
* a user-emitted `Overwrite` in `checkpoint_writes` — replaces
|
||||
prior history;
|
||||
* a non-sentinel blob in `checkpoint_blobs` — a pre-delta snapshot;
|
||||
bound as `DeltaChannelWrites.seed` so replay starts from it.
|
||||
Walk is newest → oldest from the target's parent. A non-sentinel
|
||||
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
|
||||
walk and is bound as `DeltaChannelWrites.seed` so replay starts
|
||||
from it.
|
||||
|
||||
Writes stored at `target_id` itself are pending writes for the next
|
||||
step and are excluded — the walk begins at the target's parent.
|
||||
"""
|
||||
overwrite_types = _overwrite_types()
|
||||
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_of: dict[str, str | None] = {}
|
||||
for r in parents_rows:
|
||||
@@ -298,15 +292,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
seed = blob_value
|
||||
found_seed = True
|
||||
break
|
||||
terminated = False
|
||||
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(val)
|
||||
if isinstance(val, overwrite_types):
|
||||
terminated = True
|
||||
break
|
||||
if terminated:
|
||||
break
|
||||
|
||||
collected.reverse() # oldest → newest
|
||||
if found_seed:
|
||||
|
||||
@@ -48,21 +48,6 @@ _DELTA_RECONSTRUCTION: contextvars.ContextVar[bool] = contextvars.ContextVar(
|
||||
)
|
||||
|
||||
|
||||
def _overwrite_types() -> tuple[type, ...]:
|
||||
"""Return `(Overwrite,)` if `langgraph` is installed, else `()`.
|
||||
|
||||
`Overwrite` lives in `langgraph.types`, which this library does not depend
|
||||
on; importing eagerly would also be circular. An empty tuple makes
|
||||
`isinstance(x, overwrite_types)` safely return `False` when `langgraph` is
|
||||
not installed — no `Overwrite` values can exist in that environment.
|
||||
"""
|
||||
try:
|
||||
from langgraph.types import Overwrite # type: ignore[import-untyped]
|
||||
except ImportError:
|
||||
return ()
|
||||
return (Overwrite,)
|
||||
|
||||
|
||||
def _split_list_config(
|
||||
config: RunnableConfig,
|
||||
) -> tuple[RunnableConfig, RunnableConfig | None]:
|
||||
@@ -527,8 +512,8 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
Default `SEED_UNSET` means no seed.
|
||||
|
||||
Walks the **parent chain** (not `list(before=...)`): for a thread with
|
||||
forks, only on-path ancestors contribute. Scans newest→oldest and
|
||||
stops at the first `Overwrite`, so reconstruction cost is bounded.
|
||||
forks, only on-path ancestors contribute. Writes are returned
|
||||
oldest→newest.
|
||||
|
||||
Writes stored at the target `checkpoint_id` itself are pending writes
|
||||
for the next step and are excluded — pregel applies them separately
|
||||
@@ -546,7 +531,6 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
# method ignores — it only reads pending_writes).
|
||||
if _DELTA_RECONSTRUCTION.get():
|
||||
return DeltaChannelWrites(writes=[])
|
||||
overwrite_types = _overwrite_types()
|
||||
|
||||
token = _DELTA_RECONSTRUCTION.set(True)
|
||||
try:
|
||||
@@ -566,9 +550,6 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
if ch != channel:
|
||||
continue
|
||||
collected.append(value)
|
||||
if isinstance(value, overwrite_types):
|
||||
collected.reverse()
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
cursor_config = tup.parent_config
|
||||
collected.reverse()
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
@@ -581,7 +562,6 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""Async version of `get_channel_writes`. See docstring there."""
|
||||
if _DELTA_RECONSTRUCTION.get():
|
||||
return DeltaChannelWrites(writes=[])
|
||||
overwrite_types = _overwrite_types()
|
||||
|
||||
token = _DELTA_RECONSTRUCTION.set(True)
|
||||
try:
|
||||
@@ -599,9 +579,6 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
if ch != channel:
|
||||
continue
|
||||
collected.append(value)
|
||||
if isinstance(value, overwrite_types):
|
||||
collected.reverse()
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
cursor_config = tup.parent_config
|
||||
collected.reverse()
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
|
||||
@@ -23,7 +23,6 @@ from langgraph.checkpoint.base import (
|
||||
CheckpointTuple,
|
||||
DeltaChannelWrites,
|
||||
SerializerProtocol,
|
||||
_overwrite_types,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
@@ -173,14 +172,11 @@ class InMemorySaver(
|
||||
chain.append(current)
|
||||
_, _, parent = entry
|
||||
current = parent
|
||||
overwrite_types = _overwrite_types()
|
||||
|
||||
# Scan newest→oldest. Two terminators stop the walk:
|
||||
# 1. a user-emitted `Overwrite` in writes — replaces prior history;
|
||||
# 2. a pre-delta blob on an ancestor — bind it as `seed`.
|
||||
# Without (2), a thread migrated from pre-delta storage would replay
|
||||
# ancestor writes all the way to the root AND miss any value that
|
||||
# lived only in the old blob (e.g. from `update_state`).
|
||||
# Scan newest→oldest. A pre-delta blob on an ancestor terminates the
|
||||
# walk and is bound as `seed`; without this, a thread migrated from
|
||||
# pre-delta storage would replay ancestor writes all the way to the
|
||||
# root AND miss any value that lived only in the old blob (e.g. from
|
||||
# `update_state`).
|
||||
#
|
||||
# At each ancestor, check the blob BEFORE processing its pending
|
||||
# writes: a pre-delta blob represents the state AT that ancestor,
|
||||
@@ -215,9 +211,6 @@ class InMemorySaver(
|
||||
continue
|
||||
val = self.serde.loads_typed(serialized)
|
||||
collected.append(val)
|
||||
if isinstance(val, overwrite_types):
|
||||
collected.reverse()
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
collected.reverse()
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
|
||||
|
||||
@@ -555,43 +555,6 @@ class TestBaseFallbackGetChannelWrites:
|
||||
assert results[0] == expected
|
||||
assert results[1] == expected
|
||||
|
||||
def test_fallback_stops_at_first_overwrite(self) -> None:
|
||||
"""An `Overwrite` dominates older history: scan newest→oldest stops at
|
||||
the first one (so `snapshot_every` / user Overwrites bound replay cost).
|
||||
"""
|
||||
langgraph_types = pytest.importorskip(
|
||||
"langgraph.types", reason="langgraph core not installed"
|
||||
)
|
||||
Overwrite = langgraph_types.Overwrite
|
||||
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
serde = JsonPlusSerializer()
|
||||
cp1_id = "00000000000000000000000000000002.0000000000000000"
|
||||
# Replace cp1's write with an Overwrite — cp0's write must be dropped.
|
||||
saver.writes[(thread_id, ns, cp1_id)][("task2", 0)] = (
|
||||
"task2",
|
||||
"messages",
|
||||
serde.dumps_typed(Overwrite([{"content": "reset"}])),
|
||||
"",
|
||||
)
|
||||
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver.get_channel_writes(config, "messages")
|
||||
|
||||
assert len(result.writes) == 1
|
||||
assert isinstance(result.writes[0], Overwrite)
|
||||
assert result.writes[0].value == [{"content": "reset"}]
|
||||
assert result.seed is SEED_UNSET
|
||||
|
||||
|
||||
class TestPreDeltaBlobTerminator:
|
||||
"""Verify the pre-delta blob terminator: when the ancestor walk hits a
|
||||
checkpoint whose blob for the channel is a real value (not
|
||||
|
||||
@@ -80,8 +80,6 @@ CONF = cast(Literal["configurable"], sys.intern("configurable"))
|
||||
# key for the configurable dict in RunnableConfig
|
||||
NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000")
|
||||
# the task_id to use for writes that are not associated with a task
|
||||
SNAPSHOT_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000001")
|
||||
# the task_id to use for framework-injected DeltaChannel snapshot writes
|
||||
OVERWRITE = sys.intern("__overwrite__")
|
||||
# dict key for the overwrite value, used as `{'__overwrite__': value}`
|
||||
|
||||
|
||||
@@ -11,7 +11,6 @@ from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel, Value
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.errors import EmptyChannelError
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
@@ -31,41 +30,28 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
to eliminate O(N²) blob growth — storage is O(N) using the writes table that
|
||||
every checkpointer already maintains.
|
||||
|
||||
``snapshot_every`` bounds reconstruction cost. When set, every N effective
|
||||
writes an `Overwrite(full_value)` is injected into `checkpoint_writes`; on
|
||||
reload the saver scans writes newest→oldest and stops at the first
|
||||
`Overwrite`, so replay work is bounded regardless of thread age. Any
|
||||
user-written `Overwrite` on the channel provides the same benefit for free.
|
||||
Reconstruction replays every ancestor write through the operator, so
|
||||
per-get cost scales with thread depth. Compaction for deep threads is
|
||||
a follow-up — today, use this on threads of a few hundred turns.
|
||||
|
||||
Usage::
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
|
||||
# With periodic snapshotting for long threads:
|
||||
messages: Annotated[
|
||||
list[AnyMessage],
|
||||
DeltaChannel(add_messages, snapshot_every=100),
|
||||
]
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"value",
|
||||
"operator",
|
||||
"snapshot_every",
|
||||
"_writes_since_snapshot",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Any, Any], Any],
|
||||
*,
|
||||
snapshot_every: int | None = None,
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.value: Any = []
|
||||
self.snapshot_every = snapshot_every
|
||||
self._writes_since_snapshot = 0
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
@@ -86,62 +72,50 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
return self.typ
|
||||
|
||||
def copy(self) -> Self:
|
||||
new: DeltaChannel[Value] = DeltaChannel(
|
||||
self.operator, snapshot_every=self.snapshot_every
|
||||
)
|
||||
new: DeltaChannel[Value] = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
||||
new._writes_since_snapshot = self._writes_since_snapshot
|
||||
return new
|
||||
|
||||
def _apply_write(self, value: Any, write: Any, counter: int) -> tuple[Any, int]:
|
||||
"""Apply one write to `value`; return (new_value, new_counter).
|
||||
def _apply_write(self, value: Any, write: Any) -> Any:
|
||||
"""Apply one write to `value` and return the new value.
|
||||
|
||||
An `Overwrite` resets the counter to 0; any other write increments it.
|
||||
Centralizes the Overwrite/reducer branching used by both `update` (live
|
||||
super-step) and `from_checkpoint` (ancestor replay).
|
||||
An `Overwrite` replaces the value; any other write is folded through
|
||||
the operator. Centralizes the Overwrite/reducer branching used by both
|
||||
`update` (live super-step) and `from_checkpoint` (ancestor replay).
|
||||
"""
|
||||
is_overwrite, overwrite_value = _get_overwrite(write)
|
||||
if is_overwrite:
|
||||
new_value = (
|
||||
return (
|
||||
_copy.copy(overwrite_value)
|
||||
if overwrite_value is not None
|
||||
else _empty(self.typ)
|
||||
)
|
||||
return new_value, 0
|
||||
base = _empty(self.typ) if value is MISSING else value
|
||||
return self.operator(base, write), counter + 1
|
||||
return self.operator(base, write)
|
||||
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
new: DeltaChannel[Value] = DeltaChannel(
|
||||
self.operator, snapshot_every=self.snapshot_every
|
||||
)
|
||||
new: DeltaChannel[Value] = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
new.value = _empty(new.typ)
|
||||
new._writes_since_snapshot = 0
|
||||
elif isinstance(checkpoint, DeltaChannelWrites):
|
||||
# Saver reconstructed per-step writes; replay through the operator.
|
||||
# `seed` (if set) is a pre-delta accumulated value that terminates
|
||||
# the ancestor walk on the saver side: replay starts from it
|
||||
# instead of the channel's empty value. Counter tracks writes
|
||||
# since the last Overwrite so snapshot cadence stays accurate
|
||||
# across reloads.
|
||||
# instead of the channel's empty value.
|
||||
value: Any = (
|
||||
_empty(new.typ) if checkpoint.seed is SEED_UNSET else checkpoint.seed
|
||||
)
|
||||
counter = 0
|
||||
for write in checkpoint.writes:
|
||||
value, counter = new._apply_write(value, write, counter)
|
||||
value = new._apply_write(value, write)
|
||||
new.value = value
|
||||
new._writes_since_snapshot = counter
|
||||
else:
|
||||
# Backward compat: a pre-DeltaChannel thread stored the accumulated
|
||||
# value directly (no saver-side reconstruction happened). Trust it.
|
||||
new.value = checkpoint
|
||||
new._writes_since_snapshot = 0
|
||||
return new
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
@@ -167,9 +141,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
elif seen_overwrite:
|
||||
# Post-Overwrite writes within the same super-step are dropped.
|
||||
continue
|
||||
self.value, self._writes_since_snapshot = self._apply_write(
|
||||
self.value, value, self._writes_since_snapshot
|
||||
)
|
||||
self.value = self._apply_write(self.value, value)
|
||||
return True
|
||||
|
||||
def get(self) -> Any:
|
||||
@@ -182,24 +154,3 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
return DELTA_SENTINEL
|
||||
|
||||
def should_snapshot(self) -> bool:
|
||||
"""True if enough writes have accumulated to justify a snapshot.
|
||||
|
||||
Pregel checks this after a checkpoint is saved; if true, it injects
|
||||
`snapshot_write()` into `checkpoint_writes`.
|
||||
"""
|
||||
return (
|
||||
self.snapshot_every is not None
|
||||
and self._writes_since_snapshot >= self.snapshot_every
|
||||
)
|
||||
|
||||
def snapshot_write(self) -> Overwrite:
|
||||
"""Return the write to persist and reset the counter.
|
||||
|
||||
The write is an `Overwrite(current_value)`; on replay the operator's
|
||||
`Overwrite` handling resets state to this value, and the saver's
|
||||
ancestor walk can stop here.
|
||||
"""
|
||||
self._writes_since_snapshot = 0
|
||||
return Overwrite(_copy.copy(self.value))
|
||||
|
||||
@@ -57,7 +57,6 @@ from langgraph._internal._constants import (
|
||||
NULL_TASK_ID,
|
||||
PUSH,
|
||||
RESUME,
|
||||
SNAPSHOT_TASK_ID,
|
||||
TASKS,
|
||||
)
|
||||
from langgraph._internal._replay import ReplayState
|
||||
@@ -69,7 +68,6 @@ from langgraph.callbacks import (
|
||||
GraphResumeEvent,
|
||||
)
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.untracked_value import UntrackedValue
|
||||
from langgraph.constants import TAG_HIDDEN
|
||||
from langgraph.errors import (
|
||||
@@ -947,32 +945,6 @@ class PregelLoop:
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
|
||||
},
|
||||
}
|
||||
# DeltaChannel snapshot injection: after the checkpoint save is
|
||||
# queued, any DeltaChannel that has crossed its `snapshot_every`
|
||||
# threshold emits an Overwrite write stored at the just-saved
|
||||
# checkpoint_id. On any descendant load, the saver's ancestor walk
|
||||
# encounters this Overwrite and stops, bounding replay cost.
|
||||
if self.checkpointer_put_writes is not None and self.channels:
|
||||
snapshot_writes: list[tuple[str, Any]] = []
|
||||
for ch_name, channel in self.channels.items():
|
||||
if isinstance(channel, DeltaChannel) and channel.should_snapshot():
|
||||
snapshot_writes.append((ch_name, channel.snapshot_write()))
|
||||
if snapshot_writes:
|
||||
if self.checkpointer_put_writes_accepts_task_path:
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
self.checkpoint_config,
|
||||
snapshot_writes,
|
||||
SNAPSHOT_TASK_ID,
|
||||
"",
|
||||
)
|
||||
else:
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
self.checkpoint_config,
|
||||
snapshot_writes,
|
||||
SNAPSHOT_TASK_ID,
|
||||
)
|
||||
if not exiting:
|
||||
# increment step
|
||||
self.step += 1
|
||||
|
||||
@@ -460,30 +460,6 @@ def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
|
||||
assert ch.get() == {"x": 10, "y": 20, "z": 30}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_snapshot_write_preserves_shape() -> None:
|
||||
"""snapshot_write() on a dict channel must emit Overwrite(dict), not Overwrite(list)."""
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
spec = _delta_channel_with_type(merge_dicts, dict)
|
||||
# Force snapshot_every by reaching into the spec instance.
|
||||
assert isinstance(spec, DeltaChannel)
|
||||
spec.snapshot_every = 2
|
||||
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
ch.update([{"a": 1}])
|
||||
ch.update([{"b": 2}])
|
||||
assert ch.should_snapshot()
|
||||
|
||||
w = ch.snapshot_write()
|
||||
assert isinstance(w, Overwrite)
|
||||
assert w.value == {"a": 1, "b": 2}
|
||||
assert isinstance(w.value, dict)
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
|
||||
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`.
|
||||
|
||||
@@ -605,69 +581,10 @@ def test_delta_channel_dict_reducer_backwards_compat() -> None:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# snapshot_every
|
||||
# seed / pre-delta migration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_counter_triggers() -> None:
|
||||
"""Counter hits threshold → should_snapshot() true; snapshot_write() resets it."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
ch = DeltaChannel(add_messages, snapshot_every=3).from_checkpoint(MISSING)
|
||||
assert not ch.should_snapshot()
|
||||
|
||||
ch.update([HumanMessage(content="a", id="1")])
|
||||
assert not ch.should_snapshot()
|
||||
ch.update([HumanMessage(content="b", id="2")])
|
||||
assert not ch.should_snapshot()
|
||||
ch.update([HumanMessage(content="c", id="3")])
|
||||
assert ch.should_snapshot()
|
||||
|
||||
w = ch.snapshot_write()
|
||||
assert isinstance(w, Overwrite)
|
||||
assert len(w.value) == 3
|
||||
# counter reset
|
||||
assert not ch.should_snapshot()
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_default_disabled() -> None:
|
||||
"""No snapshot_every → should_snapshot() is never true."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
for i in range(50):
|
||||
ch.update([HumanMessage(content=str(i), id=str(i))])
|
||||
assert not ch.should_snapshot()
|
||||
|
||||
|
||||
def test_delta_channel_user_overwrite_resets_counter() -> None:
|
||||
"""An Overwrite (user or framework) resets the snapshot counter."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
ch = DeltaChannel(add_messages, snapshot_every=3).from_checkpoint(MISSING)
|
||||
ch.update([HumanMessage(content="a", id="1")])
|
||||
ch.update([HumanMessage(content="b", id="2")])
|
||||
# Right before threshold — user Overwrite should reset.
|
||||
ch.update([Overwrite([HumanMessage(content="new", id="new")])])
|
||||
assert not ch.should_snapshot()
|
||||
# Need 3 more writes to trigger.
|
||||
ch.update([HumanMessage(content="c", id="c")])
|
||||
ch.update([HumanMessage(content="d", id="d")])
|
||||
ch.update([HumanMessage(content="e", id="e")])
|
||||
assert ch.should_snapshot()
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_honors_seed() -> None:
|
||||
"""DeltaChannelWrites(seed=...) starts replay from that snapshot.
|
||||
|
||||
@@ -714,117 +631,3 @@ def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_unset() -> Non
|
||||
assert unset.seed is SEED_UNSET
|
||||
|
||||
|
||||
def test_delta_channel_replay_tracks_counter_across_overwrite() -> None:
|
||||
"""Counter reloaded from writes reflects writes-since-last-Overwrite."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaChannelWrites
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
spec = DeltaChannel(add_messages, snapshot_every=3)
|
||||
writes = DeltaChannelWrites(
|
||||
[
|
||||
HumanMessage(content="a", id="1"),
|
||||
HumanMessage(content="b", id="2"),
|
||||
HumanMessage(content="c", id="3"),
|
||||
Overwrite([HumanMessage(content="reset", id="r")]),
|
||||
HumanMessage(content="d", id="4"),
|
||||
]
|
||||
)
|
||||
ch = spec.from_checkpoint(writes)
|
||||
# Post-snapshot: one write since reset.
|
||||
assert ch._writes_since_snapshot == 1
|
||||
assert not ch.should_snapshot()
|
||||
assert len(ch.get()) == 2 # reset → [r], +d → [r, d]
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_end_to_end_inmemory() -> None:
|
||||
"""Full graph: snapshot is injected, descendants short-circuit at it."""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)]
|
||||
|
||||
n = {"v": 0}
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n["v"] += 1
|
||||
return {"messages": [AIMessage(content=f"r{n['v']}", id=f"ai{n['v']}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "snap"}}
|
||||
|
||||
# 3 turns × (HumanMessage input + AIMessage reply) — plenty to trigger
|
||||
# multiple snapshots at snapshot_every=2.
|
||||
for i in range(3):
|
||||
graph.invoke({"messages": [HumanMessage(content=f"q{i}", id=f"h{i}")]}, config)
|
||||
|
||||
# Final state has 6 messages (3 H + 3 AI) regardless of snapshots.
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 6
|
||||
|
||||
# At least one Overwrite snapshot write was injected into checkpoint_writes.
|
||||
all_writes = [
|
||||
v
|
||||
for wdict in saver.writes.values()
|
||||
for (_task, _idx), (_, ch, serialized, _) in wdict.items()
|
||||
if ch == "messages"
|
||||
for v in [saver.serde.loads_typed(serialized)]
|
||||
]
|
||||
overwrites = [w for w in all_writes if isinstance(w, Overwrite)]
|
||||
assert len(overwrites) >= 1, (
|
||||
f"expected at least one snapshot Overwrite, got writes: {all_writes}"
|
||||
)
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_preserves_time_travel() -> None:
|
||||
"""Time-travel to a checkpoint created before a snapshot still replays correctly."""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"r{n}", id=f"ai{n}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "snap-tt"}}
|
||||
|
||||
for i in range(4):
|
||||
graph.invoke({"messages": [HumanMessage(content=f"q{i}", id=f"h{i}")]}, config)
|
||||
|
||||
# Walk history; each snapshot has the expected message count at that point.
|
||||
history = list(graph.get_state_history(config))
|
||||
# Reverse to chronological order.
|
||||
history = list(reversed(history))
|
||||
# At each point the visible message count should monotonically grow.
|
||||
counts = [len(h.values.get("messages", [])) for h in history]
|
||||
assert counts == sorted(counts), f"message counts not monotonic: {counts}"
|
||||
|
||||
@@ -126,15 +126,6 @@ class DeltaState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
|
||||
_SNAPSHOT_EVERY = 25
|
||||
|
||||
|
||||
class DeltaSnapshotState(TypedDict):
|
||||
messages: Annotated[
|
||||
list, DeltaChannel(add_messages, snapshot_every=_SNAPSHOT_EVERY)
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factory
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -232,9 +223,9 @@ def _approx_tokens(n_turns: int) -> str:
|
||||
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
|
||||
TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
|
||||
# Snapshot-only counts: pure delta replay becomes painful past 500 turns and
|
||||
# add_messages storage balloons past 1 GB, so only the snapshot variant runs here.
|
||||
SNAPSHOT_ONLY_TURN_COUNTS = [1000]
|
||||
# Deep-thread counts where add_messages blob storage would exceed 1 GB;
|
||||
# only DeltaChannel runs here.
|
||||
DELTA_ONLY_TURN_COUNTS = [1000]
|
||||
|
||||
|
||||
def _checkpointer_factories() -> list[tuple[str, Any]]:
|
||||
@@ -294,24 +285,22 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = []
|
||||
for turns in TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
_, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
|
||||
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
|
||||
with _make_saver() as saver:
|
||||
_, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt))
|
||||
for turns in DELTA_ONLY_TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
_, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
|
||||
rows.append((turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt))
|
||||
for turns in SNAPSHOT_ONLY_TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
_, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
|
||||
rows.append((turns, None, None, s_bytes, None, None, s_rt))
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, None, d_bytes, None, d_rt, None, d_wt))
|
||||
|
||||
# ── Table 1: Storage ─────────────────────────────────────────────────────
|
||||
W = 78
|
||||
W = 64
|
||||
print("Storage (checkpoint blob bytes)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} "
|
||||
f"{'delta+snap':>12} {'savings':>8}"
|
||||
f"{'savings':>8}"
|
||||
)
|
||||
print("-" * W)
|
||||
|
||||
@@ -325,7 +314,7 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
def _ms_or_na(v: Any) -> str:
|
||||
return "n/a" if v is None else f"{v * 1000:.1f}ms"
|
||||
|
||||
for turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt in rows:
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0:
|
||||
ratio_str = "n/a"
|
||||
else:
|
||||
@@ -334,35 +323,48 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} "
|
||||
f"{_bytes_or_na(s_bytes):>12} {ratio_str:>8}"
|
||||
f"{ratio_str:>8}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
# ── Table 2: Read latency ─────────────────────────────────────────────────
|
||||
print("Read latency (avg of 5 get_state calls = cost per invoke)")
|
||||
print("Read latency (avg of 5 get_state calls)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} "
|
||||
f"{'delta+snap':>12}"
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}"
|
||||
)
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt in rows:
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12} "
|
||||
f"{_ms_or_na(s_rt):>12}"
|
||||
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
# ── Table 3: Per-invoke latency (total write_elapsed / turns) ─────────────
|
||||
print("Per-invoke latency (total graph.invoke time / turns)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}"
|
||||
)
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
def _per(wt: Any) -> str:
|
||||
if wt is None:
|
||||
return "n/a"
|
||||
return f"{(wt / turns) * 1000:.1f}ms"
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_per(b_wt):>12} {_per(d_wt):>12}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
print("Legend:")
|
||||
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
|
||||
print(" delta = DeltaChannel(add_messages) — O(N) storage")
|
||||
print(
|
||||
f" delta+snap = DeltaChannel(add_messages, snapshot_every={_SNAPSHOT_EVERY})"
|
||||
" — O(N) storage, bounded read"
|
||||
)
|
||||
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
|
||||
print(" delta = Annotated[list, DeltaChannel(add_messages)] — O(N) storage")
|
||||
print()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user