diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index f6a0d8645..6d727b1ba 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -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: diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index f27a2e740..af65db1c7 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -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) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 05633418f..b02d16af1 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -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) diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 322a6f2d3..e32fbb5b6 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -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 diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index a7e128d79..68cb48fe8 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -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}` diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 56cc8123f..28e223a61 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -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)) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 22e41ac97..976c1382b 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -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 diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 2f049f953..de0cea7a6 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -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}" diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index c9b7a837f..bc8b9d441 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -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()