From 4a6ddbb0e30c07586244c2c05da90f5c2e7c2a27 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 22 Apr 2026 13:08:52 -0400 Subject: [PATCH] feat(delta-channel): store sentinel in blobs, reconstruct from checkpoint_writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DeltaChannel.checkpoint() now returns a zero-byte DeltaChannelSentinel instead of duplicating delta data in checkpoint_blobs. Reconstruction walks the parent checkpoint chain via checkpoint_writes (which already holds per-step writes) and replays them through the operator. In-memory benchmark (100 turns, ~20K tokens): storage: 10.2 MB → 40.5 KB (251x reduction) read: 0.6ms → 7.9ms (reconstruction cost, amortized by storage savings) InMemorySaver and PostgresSaver override get_channel_writes() with efficient implementations (Python dict walk and recursive CTE respectively). The base class fallback uses self.list() with a thread-local recursion guard. --- .../langgraph/checkpoint/postgres/__init__.py | 24 +- .../langgraph/checkpoint/postgres/aio.py | 90 +++--- .../langgraph/checkpoint/postgres/base.py | 112 +++----- .../langgraph/checkpoint/base/__init__.py | 70 ++++- .../langgraph/checkpoint/memory/__init__.py | 155 +++++------ .../langgraph/checkpoint/serde/_msgpack.py | 2 - .../langgraph/checkpoint/serde/jsonplus.py | 19 +- libs/checkpoint/tests/test_jsonplus.py | 9 +- libs/checkpoint/tests/test_memory.py | 87 +++--- libs/langgraph/langgraph/channels/delta.py | 118 ++------ libs/langgraph/tests/test_channels.py | 263 +++++------------- .../tests/test_delta_channel_benchmark.py | 64 ++--- 12 files changed, 410 insertions(+), 603 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 5d54dfb18..cd5a14a19 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -4,7 +4,7 @@ import threading from collections import defaultdict from collections.abc import Iterator, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -442,14 +442,22 @@ class PostgresSaver(BasePostgresSaver): including its configuration, metadata, parent checkpoint (if any), and pending writes. """ - with self._cursor() as cur: - channel_values = self._load_blobs( - value["channel_values"], - thread_id=value["thread_id"], - checkpoint_ns=value["checkpoint_ns"], - checkpoint_id=value["checkpoint_id"], - cur=cur, + from langgraph.checkpoint.base import DeltaChannelSentinel + + channel_values = self._load_blobs(value["channel_values"]) + if any(isinstance(v, DeltaChannelSentinel) for v in channel_values.values()): + cp_config = cast( + RunnableConfig, + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, ) + with self._cursor() as cur: + self._resolve_delta_channels(cp_config, channel_values, cur) return CheckpointTuple( { "configurable": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 1027c18fd..5ab8d5994 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -13,8 +13,7 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, - DeltaChainValue, - DeltaValue, + DeltaChannelSentinel, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -393,80 +392,55 @@ class AsyncPostgresSaver(BasePostgresSaver): async with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - async def _aload_delta_chain( + async def _aget_channel_writes_cur( self, thread_id: str, checkpoint_ns: str, checkpoint_id: str, channel: str, cur: Any, - ) -> DeltaChainValue: - """Fetch the full delta chain for a channel in one recursive CTE query (async).""" + ) -> list[Any]: + """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async).""" await cur.execute( """ - WITH RECURSIVE chain AS ( - SELECT - c.checkpoint_id, - c.parent_checkpoint_id, - cb.version, - cb.type, - cb.blob - FROM checkpoints c - JOIN checkpoint_blobs cb - ON cb.thread_id = %s - AND cb.checkpoint_ns = %s - AND cb.channel = %s - AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text - WHERE c.thread_id = %s - AND c.checkpoint_ns = %s - AND c.checkpoint_id = %s - + WITH RECURSIVE chain(cid, depth) AS ( + SELECT parent_checkpoint_id, 0 + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s UNION ALL - - SELECT - c.checkpoint_id, - c.parent_checkpoint_id, - cb.version, - cb.type, - cb.blob - FROM chain prev - JOIN checkpoints c ON c.checkpoint_id = prev.parent_checkpoint_id - JOIN checkpoint_blobs cb - ON cb.thread_id = %s - AND cb.checkpoint_ns = %s - AND cb.channel = %s - AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text - WHERE prev.parent_checkpoint_id IS NOT NULL - AND prev.type = 'delta' + SELECT c.parent_checkpoint_id, ch.depth + 1 + FROM checkpoints c + JOIN chain ch ON c.checkpoint_id = ch.cid + WHERE ch.cid IS NOT NULL ) - SELECT DISTINCT ON (version) type, blob - FROM chain - ORDER BY version ASC + SELECT cw.type, cw.blob + FROM checkpoint_writes cw + JOIN chain ON cw.checkpoint_id = chain.cid + WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s + ORDER BY chain.depth DESC, cw.task_id, cw.idx """, ( - thread_id, - checkpoint_ns, - channel, - channel, thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns, channel, - channel, ), ) rows = await cur.fetchall() - base = None - deltas: list[list[Any]] = [] - for row in rows: - blob = self.serde.loads_typed((row["type"], row["blob"])) - if isinstance(blob, DeltaValue): - deltas.append(blob.delta) - else: - base = blob - return DeltaChainValue(base=base, deltas=deltas) + return [self.serde.loads_typed((row["type"], row["blob"])) for row in rows] + + async def aget_channel_writes( + self, config: RunnableConfig, channel: str + ) -> list[Any]: + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"]["checkpoint_id"] + async with self._cursor() as cur: + return await self._aget_channel_writes_cur( + thread_id, checkpoint_ns, checkpoint_id, channel, cur + ) async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: """ @@ -489,12 +463,14 @@ class AsyncPostgresSaver(BasePostgresSaver): if blob_values: channel_values = self._load_blobs(blob_values) delta_channels = [ - k.decode() for k, t, _ in blob_values if t.decode() == "delta" + ch + for ch, v in channel_values.items() + if isinstance(v, DeltaChannelSentinel) ] if delta_channels: async with self._cursor() as cur: for channel in delta_channels: - channel_values[channel] = await self._aload_delta_chain( + channel_values[channel] = await self._aget_channel_writes_cur( thread_id, checkpoint_ns, checkpoint_id, channel, cur ) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index cf571c884..66eacb082 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -11,8 +11,7 @@ from langgraph.checkpoint.base import ( WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, - DeltaChainValue, - DeltaValue, + DeltaChannelSentinel, get_checkpoint_id, ) from langgraph.checkpoint.serde.types import TASKS @@ -188,105 +187,72 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): def _load_blobs( self, - blob_values: list[tuple[bytes, bytes, bytes]], - *, - thread_id: str = "", - checkpoint_ns: str = "", - checkpoint_id: str = "", - cur: Any = None, + blob_values: Any, ) -> dict[str, Any]: if not blob_values: return {} result: dict[str, Any] = {} - delta_channels: list[str] = [] for k, t, v in blob_values: - channel = k.decode() type_tag = t.decode() - if type_tag == "delta": - delta_channels.append(channel) - elif type_tag != "empty": - result[channel] = self.serde.loads_typed((type_tag, v)) - if delta_channels and cur is not None and checkpoint_id: - for channel in delta_channels: - result[channel] = self._load_delta_chain( - thread_id, checkpoint_ns, checkpoint_id, channel, cur - ) + if type_tag != "empty": + result[k.decode()] = self.serde.loads_typed((type_tag, v)) return result - def _load_delta_chain( + def _resolve_delta_channels( + self, + config: RunnableConfig, + channel_values: dict[str, Any], + cur: Any, + ) -> None: + for channel, value in list(channel_values.items()): + if isinstance(value, DeltaChannelSentinel): + channel_values[channel] = self._get_channel_writes_cur( + config["configurable"]["thread_id"], + config["configurable"].get("checkpoint_ns", ""), + config["configurable"]["checkpoint_id"], + channel, + cur, + ) + + def _get_channel_writes_cur( self, thread_id: str, checkpoint_ns: str, checkpoint_id: str, channel: str, cur: Any, - ) -> DeltaChainValue: - """Fetch the full delta chain for a channel in one recursive CTE query.""" + ) -> list[Any]: + """Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest.""" cur.execute( """ - WITH RECURSIVE chain AS ( - SELECT - c.checkpoint_id, - c.parent_checkpoint_id, - cb.version, - cb.type, - cb.blob - FROM checkpoints c - JOIN checkpoint_blobs cb - ON cb.thread_id = %s - AND cb.checkpoint_ns = %s - AND cb.channel = %s - AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text - WHERE c.thread_id = %s - AND c.checkpoint_ns = %s - AND c.checkpoint_id = %s - + WITH RECURSIVE chain(cid, depth) AS ( + SELECT parent_checkpoint_id, 0 + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s UNION ALL - - SELECT - c.checkpoint_id, - c.parent_checkpoint_id, - cb.version, - cb.type, - cb.blob - FROM chain prev - JOIN checkpoints c ON c.checkpoint_id = prev.parent_checkpoint_id - JOIN checkpoint_blobs cb - ON cb.thread_id = %s - AND cb.checkpoint_ns = %s - AND cb.channel = %s - AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text - WHERE prev.parent_checkpoint_id IS NOT NULL - AND prev.type = 'delta' + SELECT c.parent_checkpoint_id, ch.depth + 1 + FROM checkpoints c + JOIN chain ch ON c.checkpoint_id = ch.cid + WHERE ch.cid IS NOT NULL ) - SELECT DISTINCT ON (version) type, blob - FROM chain - ORDER BY version ASC + SELECT cw.type, cw.blob + FROM checkpoint_writes cw + JOIN chain ON cw.checkpoint_id = chain.cid + WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s + ORDER BY chain.depth DESC, cw.task_id, cw.idx """, ( - thread_id, - checkpoint_ns, - channel, - channel, thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns, channel, - channel, ), ) - rows = cur.fetchall() - base = None - deltas: list[list[Any]] = [] - for row in rows: - blob = self.serde.loads_typed((row["type"], row["blob"])) - if isinstance(blob, DeltaValue): - deltas.append(blob.delta) - else: - base = blob - return DeltaChainValue(base=base, deltas=deltas) + return [ + self.serde.loads_typed((row["type"], row["blob"])) for row in cur.fetchall() + ] def _dump_blobs( self, diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 6d70a1075..f52b9b117 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -3,10 +3,12 @@ from __future__ import annotations import copy import dataclasses import logging +import threading from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence from typing import ( # noqa: UP035 Any, Generic, + List, Literal, NamedTuple, TypedDict, @@ -32,19 +34,17 @@ PendingWrite = tuple[str, str, Any] @dataclasses.dataclass -class DeltaValue: - """Returned by DeltaChannel.checkpoint(). Represents one step's writes.""" +class DeltaChannelSentinel: + """Marker stored in checkpoint_blobs for a DeltaChannel field. - delta: list[Any] + No data is stored here — the actual per-step writes live in checkpoint_writes + and are replayed through the reducer at load time. + """ + + pass -@dataclasses.dataclass -class DeltaChainValue: - """Passed to DeltaChannel.from_checkpoint(). Assembled by the pregel layer.""" - - base: list[Any] | None # starting accumulated value; None = start from empty - deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest - +_DELTA_RECONSTRUCTION: threading.local = threading.local() logger = logging.getLogger(__name__) @@ -475,6 +475,56 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006 + """Collect all writes for `channel` across this checkpoint's ancestry, oldest→newest. + + Default implementation walks the full thread history via `list()`. Savers can + override with a more efficient query (InMemorySaver and PostgresSaver do this). + """ + # Guard against re-entrant calls: when list() triggers reconstruction which + # calls list() again, the inner call returns tuples with DeltaChannelSentinel + # in channel_values (which get_channel_writes ignores — it only reads + # pending_writes). This breaks the recursion safely. + if getattr(_DELTA_RECONSTRUCTION, "active", False): + return [] + _DELTA_RECONSTRUCTION.active = True + try: + result: list[Any] = [] + target_id = config["configurable"].get("checkpoint_id") + for tup in self.list(config): + if tup.config["configurable"].get("checkpoint_id") == target_id: + continue # skip the checkpoint itself; we want its ancestors' writes + if tup.pending_writes: + for _, ch, value in tup.pending_writes: + if ch == channel: + result.append(value) + result.reverse() # list() yields newest→oldest; we want oldest→newest + return result + finally: + _DELTA_RECONSTRUCTION.active = False + + async def aget_channel_writes( + self, config: RunnableConfig, channel: str + ) -> List[Any]: # noqa: UP006 + """Async version of get_channel_writes.""" + if getattr(_DELTA_RECONSTRUCTION, "active", False): + return [] + _DELTA_RECONSTRUCTION.active = True + try: + result: list[Any] = [] + target_id = config["configurable"].get("checkpoint_id") + async for tup in self.alist(config): + if tup.config["configurable"].get("checkpoint_id") == target_id: + continue + if tup.pending_writes: + for _, ch, value in tup.pending_writes: + if ch == channel: + result.append(value) + result.reverse() + return result + finally: + _DELTA_RECONSTRUCTION.active = False + def get_next_version(self, current: V | None, channel: None) -> V: """Generate the next version ID for a channel. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 193eca07d..76b3fe2f1 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -9,7 +9,7 @@ from collections import defaultdict from collections.abc import AsyncIterator, Iterator, Sequence from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack from types import TracebackType -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig @@ -20,8 +20,7 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, - DeltaChainValue, - DeltaValue, + DeltaChannelSentinel, SerializerProtocol, get_checkpoint_id, get_checkpoint_metadata, @@ -127,68 +126,56 @@ class InMemorySaver( thread_id: str, checkpoint_ns: str, versions: ChannelVersions, - checkpoint_id: str = "", ) -> dict[str, Any]: - channel_values: dict[str, Any] = {} - delta_channels: list[str] = [] - for k, v in versions.items(): - kk = (thread_id, checkpoint_ns, k, v) + result: dict[str, Any] = {} + for k, ver in versions.items(): + kk = (thread_id, checkpoint_ns, k, ver) if kk not in self.blobs: continue vv = self.blobs[kk] - if vv[0] == "delta": - delta_channels.append(k) - elif vv[0] != "empty": - channel_values[k] = self.serde.loads_typed(vv) - for channel in delta_channels: - channel_values[channel] = self._assemble_delta_chain( - thread_id, checkpoint_ns, checkpoint_id, channel - ) - return channel_values + if vv[0] == "empty": + continue + result[k] = self.serde.loads_typed(vv) + return result - def _assemble_delta_chain( + def _resolve_delta_channels( self, - thread_id: str, - checkpoint_ns: str, - checkpoint_id: str, - channel: str, - ) -> DeltaChainValue: - """Walk the checkpoint parent tree to collect all delta blobs for a channel.""" + config: RunnableConfig, + channel_values: dict[str, Any], + ) -> None: + """Replace DeltaChannelSentinel entries with reconstructed write lists.""" + for channel, value in list(channel_values.items()): + if isinstance(value, DeltaChannelSentinel): + channel_values[channel] = self.get_channel_writes(config, channel) + + def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]: + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id", "") ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {}) - blobs: list[Any] = [] - current_id: str | None = checkpoint_id - seen_versions: set[str] = set() - while current_id is not None: - entry = ns_storage.get(current_id) + # Walk the parent chain, collecting checkpoint IDs oldest→newest. + chain: list[str] = [] + current: str | None = checkpoint_id + while current is not None: + entry = ns_storage.get(current) if entry is None: break - checkpoint = self.serde.loads_typed(entry[0]) - version = checkpoint["channel_versions"].get(channel) - if version is None: - break # channel not yet in this checkpoint - _, _, current_id = entry # advance to parent before the continue/break - if version in seen_versions: - continue # same blob already collected; keep walking to older checkpoints - seen_versions.add(version) - kk = (thread_id, checkpoint_ns, channel, version) - if kk not in self.blobs: - break - vv = self.blobs[kk] - if vv[0] == "empty": - break - blob = self.serde.loads_typed(vv) - blobs.append(blob) - if not isinstance(blob, DeltaValue): - break # hit a snapshot (plain list) — chain root found - blobs.reverse() - base = None - deltas: list[list[Any]] = [] - for blob in blobs: - if isinstance(blob, DeltaValue): - deltas.append(blob.delta) - else: - base = blob - return DeltaChainValue(base=base, deltas=deltas) + chain.append(current) + _, _, parent = entry + current = parent + # Collect writes for `channel` from each checkpoint in oldest→newest order. + result: list[Any] = [] + for cp_id in reversed(chain): + step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) + for (_task_id, _idx), (_, ch, serialized, _) in sorted(step_writes.items()): + if ch == channel: + result.append(self.serde.loads_typed(serialized)) + return result + + async def aget_channel_writes( + self, config: RunnableConfig, channel: str + ) -> list[Any]: + return self.get_channel_writes(config, channel) def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the in-memory storage. @@ -211,16 +198,17 @@ class InMemorySaver( checkpoint, metadata, parent_checkpoint_id = saved writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) + channel_values = self._load_blobs( + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + ) + self._resolve_delta_channels(config, channel_values) return CheckpointTuple( config=config, checkpoint={ **checkpoint_, - "channel_values": self._load_blobs( - thread_id, - checkpoint_ns, - checkpoint_["channel_versions"], - checkpoint_id, - ), + "channel_values": channel_values, }, metadata=self.serde.loads_typed(metadata), pending_writes=[ @@ -244,22 +232,27 @@ class InMemorySaver( checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id] writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() checkpoint_ = self.serde.loads_typed(checkpoint) - return CheckpointTuple( - config={ + resolved_config = cast( + RunnableConfig, + { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": checkpoint_id, } }, + ) + channel_values = self._load_blobs( + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + ) + self._resolve_delta_channels(resolved_config, channel_values) + return CheckpointTuple( + config=resolved_config, checkpoint={ **checkpoint_, - "channel_values": self._load_blobs( - thread_id, - checkpoint_ns, - checkpoint_["channel_versions"], - checkpoint_id, - ), + "channel_values": channel_values, }, metadata=self.serde.loads_typed(metadata), pending_writes=[ @@ -354,22 +347,28 @@ class InMemorySaver( checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) - yield CheckpointTuple( - config={ + list_config = cast( + RunnableConfig, + { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": checkpoint_id, } }, + ) + channel_values = self._load_blobs( + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + ) + self._resolve_delta_channels(list_config, channel_values) + + yield CheckpointTuple( + config=list_config, checkpoint={ **checkpoint_, - "channel_values": self._load_blobs( - thread_id, - checkpoint_ns, - checkpoint_["channel_versions"], - checkpoint_id, - ), + "channel_values": channel_values, }, metadata=metadata, parent_config=( diff --git a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py index 2234fb4e8..35ac48665 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py @@ -80,8 +80,6 @@ SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset( ("langgraph.types", "Overwrite"), ("langgraph.store.base", "Item"), ("langgraph.store.base", "GetOp"), - # DeltaChannel checkpoint value type - ("langgraph.checkpoint.base", "DeltaValue"), } ) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 69180db4f..a3bb23b59 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -64,10 +64,12 @@ def _warn_once( logger.warning(msg, *args) -def _is_delta_value(obj: Any) -> bool: - from langgraph.checkpoint.base import DeltaValue # lazy import avoids circular dep +def _get_delta_sentinel_cls() -> type: + from langgraph.checkpoint.base import ( + DeltaChannelSentinel, + ) # lazy import avoids circular dep - return isinstance(obj, DeltaValue) + return DeltaChannelSentinel class JsonPlusSerializer(SerializerProtocol): @@ -262,8 +264,8 @@ class JsonPlusSerializer(SerializerProtocol): return "bytes", obj elif isinstance(obj, bytearray): return "bytearray", obj - elif _is_delta_value(obj): - return "delta", _msgpack_enc({"d": obj.delta}) + elif isinstance(obj, _get_delta_sentinel_cls()): + return "delta", b"" else: try: return "msgpack", _msgpack_enc(obj) @@ -287,12 +289,9 @@ class JsonPlusSerializer(SerializerProtocol): data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS ) elif type_ == "delta": - from langgraph.checkpoint.base import DeltaValue # lazy import + from langgraph.checkpoint.base import DeltaChannelSentinel - raw = ormsgpack.unpackb( - data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS - ) - return DeltaValue(delta=raw["d"]) + return DeltaChannelSentinel() elif self.pickle_fallback and type_ == "pickle": return pickle.loads(data_) else: diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index 363fe5aca..17a9a98b1 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -999,14 +999,13 @@ def test_msgpack_nested_pydantic_serializes_as_dict( assert result == obj -def test_delta_value_serde_round_trip() -> None: - from langgraph.checkpoint.base import DeltaValue +def test_delta_channel_sentinel_serde_round_trip() -> None: + from langgraph.checkpoint.base import DeltaChannelSentinel from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer serde = JsonPlusSerializer() - original = DeltaValue(delta=[{"type": "human", "content": "hi"}]) + original = DeltaChannelSentinel() type_tag, blob = serde.dumps_typed(original) assert type_tag == "delta" loaded = serde.loads_typed((type_tag, blob)) - assert isinstance(loaded, DeltaValue) - assert loaded.delta == original.delta + assert isinstance(loaded, DeltaChannelSentinel) diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 5c9ca108d..8c636bcb5 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -323,11 +323,10 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None: class TestInMemorySaverDeltaChannel: - def test_load_blobs_assembles_delta_chain(self) -> None: - """_load_blobs returns DeltaChainValue for delta channels, not raw DeltaValue.""" + def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None: + """_load_blobs returns DeltaChannelSentinel for delta channels (reconstruction deferred).""" from langgraph.checkpoint.base import ( - DeltaChainValue, - DeltaValue, + DeltaChannelSentinel, empty_checkpoint, ) @@ -336,55 +335,59 @@ class TestInMemorySaverDeltaChannel: thread_id, ns, channel = "t1", "", "messages" v1 = "00000000000000000000000000000001.0000000000000000" - v2 = "00000000000000000000000000000002.0000000000000000" - delta1 = DeltaValue(delta=[{"content": "hi"}]) - delta2 = DeltaValue(delta=[{"content": "bye"}]) - saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(delta1) - saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(delta2) + sentinel = DeltaChannelSentinel() + saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(sentinel) cp1 = empty_checkpoint() cp1["id"] = "cp1" cp1["channel_versions"][channel] = v1 + saver.storage[thread_id][ns] = { + "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), + } + + result = saver._load_blobs(thread_id, ns, {channel: v1}) + assert channel in result + assert isinstance(result[channel], DeltaChannelSentinel) + + def test_get_channel_writes_collects_writes(self) -> None: + """get_channel_writes collects per-step writes oldest→newest.""" + from langgraph.checkpoint.base import empty_checkpoint + + saver = InMemorySaver() + serde = JsonPlusSerializer() + + thread_id, ns, channel = "t1", "", "messages" + + cp1 = empty_checkpoint() + cp1["id"] = "cp1" cp2 = empty_checkpoint() cp2["id"] = "cp2" - cp2["channel_versions"][channel] = v2 saver.storage[thread_id][ns] = { "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), "cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"), } - - result = saver._load_blobs(thread_id, ns, {channel: v2}, "cp2") - assert channel in result - chain = result[channel] - assert isinstance(chain, DeltaChainValue) - assert chain.deltas == [[{"content": "hi"}], [{"content": "bye"}]] - - def test_load_blobs_single_delta_no_parent(self) -> None: - """Single delta with no parent checkpoint produces a chain with one delta.""" - from langgraph.checkpoint.base import ( - DeltaChainValue, - DeltaValue, - empty_checkpoint, + # cp1 has a write for channel + saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = ( + "task1", + channel, + serde.dumps_typed({"content": "hi"}), + "", + ) + # cp2 has a write for channel + saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = ( + "task2", + channel, + serde.dumps_typed({"content": "bye"}), + "", ) - saver = InMemorySaver() - serde = JsonPlusSerializer() - - thread_id, ns, channel = "t1", "", "messages" - v1 = "00000000000000000000000000000001.0000000000000000" - delta = DeltaValue(delta=[{"content": "only"}]) - saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(delta) - - cp1 = empty_checkpoint() - cp1["id"] = "cp1" - cp1["channel_versions"][channel] = v1 - saver.storage[thread_id][ns] = { - "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None) + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": "cp2", + } } - - result = saver._load_blobs(thread_id, ns, {channel: v1}, "cp1") - chain = result[channel] - assert isinstance(chain, DeltaChainValue) - assert chain.base is None - assert chain.deltas == [[{"content": "only"}]] + result = saver.get_channel_writes(config, channel) + assert result == [{"content": "hi"}, {"content": "bye"}] diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 07618668b..fb8b935c6 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Sequence from typing import Any, Generic -from langgraph.checkpoint.base import DeltaChainValue, DeltaValue +from langgraph.checkpoint.base import DeltaChannelSentinel from typing_extensions import Self from langgraph._internal._typing import MISSING @@ -14,64 +14,43 @@ from langgraph.errors import EmptyChannelError __all__ = ("DeltaChannel",) -class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): - """A channel that stores only per-step write deltas in checkpoints. +class DeltaChannel( + Generic[Value], BaseChannel[list[Value], Value, DeltaChannelSentinel] +): + """A channel that stores only a sentinel in checkpoints; per-step writes are + stored in checkpoint_writes and replayed through the operator at load time. - Reconstructs the full accumulated list at load time by replaying the - chain of deltas through the operator. Use with append-style reducers - (e.g. `add_messages`) on long-running threads to reduce checkpoint - storage from O(N²) to O(N). + Use with append-style reducers (e.g. `add_messages`) on long-running threads + to eliminate O(N²) blob growth — storage is O(N) using the writes table that + every checkpointer already maintains. - Works with all checkpointers. Savers with a dedicated blob store - (InMemorySaver, PostgresSaver) use an O(1) fast-path per chain step; - all others (SQLite, MongoDB, etc.) fall back to get_tuple traversal. - - Use `snapshot_every=N` to cap chain traversal depth at N steps. Every N - steps a full snapshot is written as the chain root; subsequent deltas - chain back to it, so `get_state` / reload never traverses more than N - checkpoints regardless of thread length. Recommended for savers without - a dedicated blob store. + Works with all checkpointers. Savers with dedicated implementations + (InMemorySaver, PostgresSaver) reconstruct in one pass; others fall back to + walking the checkpoint list. Usage:: class State(TypedDict): messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)] - # Cap reconstruction depth (recommended for SQLite / MongoDB savers): - messages: Annotated[list[AnyMessage], DeltaChannel(add_messages, snapshot_every=50)] # Dict-type reducer (type inferred from the Annotated outer type): files: Annotated[dict, DeltaChannel(merge_files)] """ - __slots__ = ( - "value", - "operator", - "snapshot_every", - "_pending", - "_base_version", - "_overwritten", - "_steps_since_snapshot", - ) + __slots__ = ("value", "operator") def __init__( self, operator: Callable[[list[Value], Any], list[Value]], *, - snapshot_every: int | None = None, + snapshot_every: int | None = None, # reserved for future use ) -> None: super().__init__(list) self.operator = operator - self.snapshot_every = snapshot_every self.value: list[Value] = [] - self._pending: list[Any] = [] - self._base_version: str | None = None - self._overwritten: bool = False - self._steps_since_snapshot: int = 0 def __eq__(self, other: object) -> bool: if not isinstance(other, DeltaChannel): return False - if self.snapshot_every != other.snapshot_every: - return False if ( self.operator.__name__ != "" and other.operator.__name__ != "" @@ -88,18 +67,14 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): return self.typ | list[self.typ] # type: ignore[name-defined] def copy(self) -> Self: - new = DeltaChannel(self.operator, snapshot_every=self.snapshot_every) + new = DeltaChannel(self.operator) new.typ = self.typ new.key = self.key new.value = self.value if self.value is MISSING else self.value.copy() - new._pending = self._pending[:] - new._base_version = self._base_version - new._overwritten = self._overwritten - new._steps_since_snapshot = self._steps_since_snapshot return new def from_checkpoint(self, checkpoint: Any) -> Self: - new = DeltaChannel(self.operator, snapshot_every=self.snapshot_every) + new = DeltaChannel(self.operator) new.typ = self.typ new.key = self.key if checkpoint is MISSING: @@ -107,29 +82,18 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): new.value = new.typ() except Exception: new.value = [] - elif isinstance(checkpoint, DeltaChainValue): - accumulated: list[Value] = ( - checkpoint.base if checkpoint.base is not None else new.typ() - ) - for step_writes in checkpoint.deltas: - for write in step_writes: - accumulated = new.operator(accumulated, write) - new.value = accumulated - # Seed the counter from actual chain depth so rehydration fires at - # the right time regardless of how many prior invocations there were. - new._steps_since_snapshot = len(checkpoint.deltas) - elif isinstance(checkpoint, DeltaValue): - raise ValueError( - f"Channel '{self.key}' uses DeltaChannel but the checkpointer " - "does not support incremental channel storage. " - "Use InMemorySaver or PostgresSaver, or remove DeltaChannel from your schema." - ) + elif isinstance(checkpoint, list): + # Flat list of individual write values (oldest→newest) from get_channel_writes. + value: Any = new.typ() + for write in checkpoint: + value = new.operator(value, write) + new.value = value else: - # Backwards compat: plain list from old BinaryOperatorAggregate checkpoint. - new.value = list(checkpoint) - new._pending = [] - new._base_version = None # set by the subsequent after_checkpoint() call - new._overwritten = False + # Backward compat: plain accumulated value (e.g. from a migrated thread). + try: + new.value = list(checkpoint) + except Exception: + new.value = [] return new def update(self, values: Sequence[Any]) -> bool: @@ -154,13 +118,10 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): self.value = ( list(overwrite_value) if overwrite_value is not None else self.typ() ) - self._pending = list(self.value) - self._overwritten = True seen_overwrite = True elif not seen_overwrite: base = self.typ() if self.value is MISSING else self.value self.value = self.operator(base, value) - self._pending.append(value) return True def get(self) -> list[Value]: @@ -171,26 +132,5 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): def is_available(self) -> bool: return self.value is not MISSING - def checkpoint(self) -> Any: - if ( - self.snapshot_every is not None - and self._steps_since_snapshot >= self.snapshot_every - ): - # Emit a full snapshot to cap chain depth at snapshot_every. - # The saver stores this as a plain (non-diff) blob, so future - # deltas will chain back to it and traversal depth resets to 1. - return list(self.value) - return DeltaValue(delta=self._pending[:]) - - def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None: - if version != self._base_version: - if self._base_version is None: - pass # First call after from_checkpoint — anchor without counting a step. - elif self.snapshot_every is not None: - if self._steps_since_snapshot >= self.snapshot_every: - self._steps_since_snapshot = 0 - else: - self._steps_since_snapshot += 1 - self._base_version = version - self._pending = [] - self._overwritten = False + def checkpoint(self) -> DeltaChannelSentinel: + return DeltaChannelSentinel() diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 74f9d5bf8..3b4ebb3c2 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -121,26 +121,22 @@ def test_untracked_value() -> None: def test_delta_channel_basic_two_steps() -> None: from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.base import DeltaValue + from langgraph.checkpoint.base import DeltaChannelSentinel from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages ch = DeltaChannel(add_messages).from_checkpoint(MISSING) - ch.after_checkpoint(None) # Step 1: one message added ch.update([HumanMessage(content="hi", id="h1")]) d1 = ch.checkpoint() - assert isinstance(d1, DeltaValue) - assert len(d1.delta) == 1 - ch.after_checkpoint("v1", checkpoint_id="cid1") + assert isinstance(d1, DeltaChannelSentinel) # Step 2: another message ch.update([AIMessage(content="hello", id="a1")]) d2 = ch.checkpoint() - assert len(d2.delta) == 1 - ch.after_checkpoint("v2") + assert isinstance(d2, DeltaChannelSentinel) # Full accumulated value is preserved in memory assert len(ch.get()) == 2 @@ -148,40 +144,21 @@ def test_delta_channel_basic_two_steps() -> None: assert ch.get()[1].content == "hello" -def test_delta_channel_after_checkpoint_no_op_when_unchanged() -> None: - 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) - ch.after_checkpoint(None) - ch.update([HumanMessage(content="hi", id="h1")]) - ch.after_checkpoint("v1") - - # Same version: no-op - ch.after_checkpoint("v1") - assert ch._base_version == "v1" - assert ch._pending == [] - - -def test_delta_channel_from_checkpoint_chain() -> None: +def test_delta_channel_from_checkpoint_writes_list() -> None: + """from_checkpoint with a flat list of individual writes replays them through the operator.""" from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.base import DeltaChainValue from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) - chain = DeltaChainValue( - base=None, - deltas=[ - [HumanMessage(content="hi", id="h1")], - [AIMessage(content="hello", id="a1")], - [HumanMessage(content="bye", id="h2")], - ], - ) - ch = spec.from_checkpoint(chain) + # Each element is one write value (as stored in checkpoint_writes) + writes = [ + HumanMessage(content="hi", id="h1"), + AIMessage(content="hello", id="a1"), + HumanMessage(content="bye", id="h2"), + ] + ch = spec.from_checkpoint(writes) msgs = ch.get() assert len(msgs) == 3 assert msgs[0].content == "hi" @@ -195,68 +172,45 @@ def test_delta_channel_from_checkpoint_backwards_compat() -> None: from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages - # Old BinaryOperatorAggregate checkpoint: plain list + # Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat spec = DeltaChannel(add_messages) old_value = [HumanMessage(content="old", id="h1")] ch = spec.from_checkpoint(old_value) assert ch.get() == old_value -def test_delta_channel_overwrite_resets_chain() -> None: +def test_delta_channel_overwrite() -> None: from langchain_core.messages import HumanMessage - from langgraph.checkpoint.base import DeltaValue + from langgraph.checkpoint.base import DeltaChannelSentinel from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages from langgraph.types import Overwrite ch = DeltaChannel(add_messages).from_checkpoint(MISSING) - ch.after_checkpoint(None) ch.update([HumanMessage(content="old", id="h1")]) - ch.after_checkpoint("v1") ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) d = ch.checkpoint() - assert isinstance(d, DeltaValue) - assert len(d.delta) == 1 - assert d.delta[0].content == "new" - # _overwritten flag must be set so next checkpoint acts as a chain root - assert ch._overwritten is True + assert isinstance(d, DeltaChannelSentinel) + # After overwrite, value is reset to only the new message + assert len(ch.get()) == 1 + assert ch.get()[0].content == "new" -def test_delta_channel_unsupported_saver_raises() -> None: - """from_checkpoint raises ValueError when the saver returns a raw DeltaValue.""" - from langgraph.checkpoint.base import DeltaValue - - from langgraph.channels.delta import DeltaChannel - from langgraph.graph.message import add_messages - - spec = DeltaChannel(add_messages) - raw = DeltaValue(delta=[{"type": "human", "content": "hello"}]) - with pytest.raises( - ValueError, match="does not support incremental channel storage" - ): - spec.from_checkpoint(raw) - - -def test_delta_channel_remove_message_delta_and_replay() -> None: - """RemoveMessage stored in a delta must round-trip correctly through the chain.""" +def test_delta_channel_remove_message_and_replay() -> None: + """RemoveMessage must round-trip correctly when writes are replayed.""" from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - from langgraph.checkpoint.base import DeltaChainValue, DeltaValue from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) ch = spec.from_checkpoint(MISSING) - ch.after_checkpoint(None) # Step 1: add two messages ch.update([HumanMessage(content="hi", id="h1")]) ch.update([AIMessage(content="hello", id="a1")]) - d1 = ch.checkpoint() - assert isinstance(d1, DeltaValue) - ch.after_checkpoint("v1", checkpoint_id="cid1") assert ch.get() == [ HumanMessage(content="hi", id="h1"), AIMessage(content="hello", id="a1"), @@ -264,127 +218,63 @@ def test_delta_channel_remove_message_delta_and_replay() -> None: # Step 2: remove the AI message ch.update([RemoveMessage(id="a1")]) - d2 = ch.checkpoint() - assert isinstance(d2, DeltaValue) - assert any(isinstance(w, RemoveMessage) for w in d2.delta) - ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == [HumanMessage(content="hi", id="h1")] - # Replay the full chain from scratch — must reproduce the post-remove state - chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta]) - ch2 = spec.from_checkpoint(chain) + # Replay the writes list from scratch — must reproduce the post-remove state + writes = [ + HumanMessage(content="hi", id="h1"), + AIMessage(content="hello", id="a1"), + RemoveMessage(id="a1"), + ] + ch2 = spec.from_checkpoint(writes) assert ch2.get() == [HumanMessage(content="hi", id="h1")] -def test_delta_channel_update_by_id_delta_and_replay() -> None: - """Updating a message by ID stored in a delta must round-trip correctly.""" +def test_delta_channel_update_by_id_and_replay() -> None: + """Updating a message by ID must round-trip correctly through writes replay.""" from langchain_core.messages import HumanMessage - from langgraph.checkpoint.base import DeltaChainValue, DeltaValue from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) ch = spec.from_checkpoint(MISSING) - ch.after_checkpoint(None) # Step 1: add a message ch.update([HumanMessage(content="original", id="h1")]) - d1 = ch.checkpoint() - assert isinstance(d1, DeltaValue) - ch.after_checkpoint("v1", checkpoint_id="cid1") # Step 2: update the same message by ID ch.update([HumanMessage(content="updated", id="h1")]) - d2 = ch.checkpoint() - assert isinstance(d2, DeltaValue) - ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == [HumanMessage(content="updated", id="h1")] - # Replay the full chain — must produce the updated message, not the original - chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta]) - ch2 = spec.from_checkpoint(chain) + # Replay writes — must produce the updated message, not the original + writes = [ + HumanMessage(content="original", id="h1"), + HumanMessage(content="updated", id="h1"), + ] + ch2 = spec.from_checkpoint(writes) assert len(ch2.get()) == 1 assert ch2.get()[0].content == "updated" -def test_delta_channel_snapshot_every_emits_plain_list() -> None: - """snapshot_every=N causes a plain-list snapshot after N steps; next deltas chain to it.""" +def test_delta_channel_checkpoint_returns_sentinel() -> None: + """checkpoint() always returns DeltaChannelSentinel regardless of state.""" + from langgraph.checkpoint.base import DeltaChannelSentinel + + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.message import add_messages + + ch = DeltaChannel(add_messages).from_checkpoint(MISSING) + assert isinstance(ch.checkpoint(), DeltaChannelSentinel) + from langchain_core.messages import HumanMessage - from langgraph.checkpoint.base import DeltaValue - from langgraph.channels.delta import DeltaChannel - from langgraph.graph.message import add_messages - - SNAP = 3 - spec = DeltaChannel(add_messages, snapshot_every=SNAP) - ch = spec.from_checkpoint(MISSING) - # First after_checkpoint anchors _base_version without counting a step. - ch.after_checkpoint("v0", checkpoint_id="cid0") - - # Steps 1..SNAP: each should stay as DeltaValue; counter increments each step. - for i in range(1, SNAP + 1): - ch.update([HumanMessage(content=f"m{i}", id=f"h{i}")]) - ckpt = ch.checkpoint() - assert isinstance(ckpt, DeltaValue), f"expected DeltaValue at step {i}" - ch.after_checkpoint(f"v{i}", checkpoint_id=f"cid{i}") - - # Step SNAP+1: _steps_since_snapshot == SNAP → snapshot fires - ch.update([HumanMessage(content="snap", id="hsnap")]) - snap = ch.checkpoint() - assert isinstance(snap, list), "expected plain-list snapshot at snapshot_every step" - assert len(snap) == SNAP + 1 - - # After snapshot, counter resets — next step is DeltaValue again - ch.after_checkpoint("vsnap", checkpoint_id="cidsnap") - ch.update([HumanMessage(content="post", id="hpost")]) - post = ch.checkpoint() - assert isinstance(post, DeltaValue) + ch.update([HumanMessage(content="hi", id="h1")]) + assert isinstance(ch.checkpoint(), DeltaChannelSentinel) -def test_delta_channel_snapshot_every_end_to_end() -> None: - """Graph with snapshot_every: get_state returns correct accumulated value after snapshot.""" - 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)] - - counter = {"n": 0} - - def respond(state: State) -> dict: - counter["n"] += 1 - return { - "messages": [ - AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}") - ] - } - - builder = StateGraph(State) - builder.add_node("respond", respond) - builder.add_edge(START, "respond") - graph = builder.compile(checkpointer=InMemorySaver()) - config = {"configurable": {"thread_id": "snap-test"}} - - # Run 5 turns — snapshot fires after 2 steps, then again after 2 more - for i in range(5): - graph.invoke({"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config) - - state = graph.get_state(config) - msgs = state.values["messages"] - # 5 human + 5 AI = 10 total - assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}" - - -def test_delta_channel_inmemory_saver_assembles_chain() -> None: - """InMemorySaver assembles the delta chain inside get_tuple (no pregel involvement).""" +def test_delta_channel_inmemory_saver_assembles_writes() -> None: + """InMemorySaver assembles writes from checkpoint_writes inside get_tuple.""" from typing import Annotated from langchain_core.messages import AIMessage, HumanMessage @@ -414,14 +304,16 @@ def test_delta_channel_inmemory_saver_assembles_chain() -> None: graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config) graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config) - # get_tuple must return a fully assembled DeltaChainValue, not raw DeltaValue - from langgraph.checkpoint.base import DeltaChainValue, DeltaValue + # get_tuple must return a resolved list (not DeltaChannelSentinel) + from langgraph.checkpoint.base import DeltaChannelSentinel saved = saver.get_tuple(config) assert saved is not None assert "messages" in saved.checkpoint["channel_values"] - assert not isinstance(saved.checkpoint["channel_values"]["messages"], DeltaValue) - assert isinstance(saved.checkpoint["channel_values"]["messages"], DeltaChainValue) + assert not isinstance( + saved.checkpoint["channel_values"]["messages"], DeltaChannelSentinel + ) + assert isinstance(saved.checkpoint["channel_values"]["messages"], list) state = graph.get_state(config) assert len(state.values["messages"]) == 4 # 2 human + 2 AI @@ -451,48 +343,39 @@ def test_delta_channel_dict_reducer_fresh_channel() -> None: def test_delta_channel_dict_reducer_basic_updates() -> None: """DeltaChannel with a dict reducer accumulates key/value pairs across steps.""" - from langgraph.checkpoint.base import DeltaValue + from langgraph.checkpoint.base import DeltaChannelSentinel def merge_dicts(left: dict, right: dict) -> dict: return {**left, **right} ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) - ch.after_checkpoint(None) ch.update([{"a": 1}]) d1 = ch.checkpoint() - assert isinstance(d1, DeltaValue) - assert d1.delta == [{"a": 1}] - ch.after_checkpoint("v1", checkpoint_id="cid1") + assert isinstance(d1, DeltaChannelSentinel) ch.update([{"b": 2}]) d2 = ch.checkpoint() - assert d2.delta == [{"b": 2}] - ch.after_checkpoint("v2") + assert isinstance(d2, DeltaChannelSentinel) assert ch.get() == {"a": 1, "b": 2} -def test_delta_channel_dict_reducer_chain_reconstruction() -> None: - """DeltaChainValue replays correctly through a dict merge reducer.""" - from langgraph.checkpoint.base import DeltaChainValue +def test_delta_channel_dict_reducer_writes_reconstruction() -> None: + """from_checkpoint with a writes list replays correctly through a dict merge reducer.""" def merge_dicts(left: dict, right: dict) -> dict: return {**left, **right} spec = _delta_channel_with_type(merge_dicts, dict) - chain = DeltaChainValue( - base={"a": 1}, - deltas=[[{"b": 2}], [{"c": 3}]], - ) - ch = spec.from_checkpoint(chain) + # Each element is one write value (oldest→newest) + writes = [{"a": 1}, {"b": 2}, {"c": 3}] + ch = spec.from_checkpoint(writes) assert ch.get() == {"a": 1, "b": 2, "c": 3} - assert ch._steps_since_snapshot == 2 def test_delta_channel_dict_reducer_with_deletions() -> None: """Dict reducer that treats None values as deletions works end-to-end (deepagents pattern).""" - from langgraph.checkpoint.base import DeltaChainValue def merge_files(left: dict | None, right: dict) -> dict: if left is None: @@ -506,25 +389,19 @@ def test_delta_channel_dict_reducer_with_deletions() -> None: return result ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING) - ch.after_checkpoint(None) ch.update([{"file1.py": "content1", "file2.py": "content2"}]) - ch.after_checkpoint("v1", checkpoint_id="cid1") # Delete file1, add file3 ch.update([{"file1.py": None, "file3.py": "content3"}]) - ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == {"file2.py": "content2", "file3.py": "content3"} - # Confirm chain reconstruction produces the same result - chain = DeltaChainValue( - base={}, - deltas=[ - [{"file1.py": "content1", "file2.py": "content2"}], - [{"file1.py": None, "file3.py": "content3"}], - ], - ) + # Confirm writes reconstruction produces the same result + writes = [ + {"file1.py": "content1", "file2.py": "content2"}, + {"file1.py": None, "file3.py": "content3"}, + ] spec = _delta_channel_with_type(merge_files, dict) - ch2 = spec.from_checkpoint(chain) + ch2 = spec.from_checkpoint(writes) assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"} diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index c9fa338a3..494a83708 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -8,6 +8,11 @@ Simulates realistic multi-turn conversations with paragraph-length messages Token estimates: 1 token ≈ 4 chars; each turn ≈ 200 tokens (human + AI). A 1M-token conversation ≈ 5,000 turns of realistic messages. + +DeltaChannel stores only a zero-byte sentinel in checkpoint_blobs; the actual +write data lives in checkpoint_writes (already stored there). Reconstruction +walks the parent chain and replays writes through the operator — O(N) total +storage vs O(N²) for plain add_messages. """ from __future__ import annotations @@ -42,8 +47,6 @@ try: except ImportError: _POSTGRES_AVAILABLE = False -SNAPSHOT_EVERY = 50 - # --------------------------------------------------------------------------- # Realistic message payload (~100 tokens / ~400 chars each) # --------------------------------------------------------------------------- @@ -121,10 +124,6 @@ class DeltaState(TypedDict): messages: Annotated[list, DeltaChannel(add_messages)] -class DeltaSnapshotState(TypedDict): - messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=SNAPSHOT_EVERY)] - - # --------------------------------------------------------------------------- # Graph factory # --------------------------------------------------------------------------- @@ -239,7 +238,12 @@ def run_benchmark() -> None: checkpointers: list[tuple[str, Any]] = [("InMemory", None)] if _POSTGRES_AVAILABLE: - checkpointers.append(("Postgres (recursive CTE)", "postgres")) + try: + import psycopg + psycopg.connect(_POSTGRES_URI).close() + checkpointers.append(("Postgres (recursive CTE)", "postgres")) + except Exception: + pass for cp_label, cp_hint in checkpointers: print(f"--- Checkpointer: {cp_label} ---") @@ -277,30 +281,28 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver) with _make_saver() as saver: d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver) - with _make_saver() as saver: - s_wt, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver) - rows.append((turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt)) + rows.append((turns, b_bytes, d_bytes, b_rt, d_rt)) # ── Table 1: Storage ───────────────────────────────────────────────────── - W = 80 + W = 70 print("Storage (checkpoint blob bytes)") print("=" * W) print( - f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'delta+snap':>12} {'savings':>8}" + f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'savings':>8}" ) print("-" * W) storage_results = [] - for turns, b_bytes, d_bytes, s_bytes, *_ in rows: + for turns, b_bytes, d_bytes, *_ in rows: if b_bytes < 0: print( - f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>12} {'n/a':>8}" + f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>8}" ) else: - ratio = b_bytes / s_bytes if s_bytes else float("inf") - storage_results.append((turns, b_bytes, s_bytes, ratio)) + ratio = b_bytes / d_bytes if d_bytes else float("inf") + storage_results.append((turns, b_bytes, d_bytes, ratio)) print( f"{turns:>6} {_approx_tokens(turns):>10} " - f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} {_fmt_bytes(s_bytes):>12} " + f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} " f"{ratio:>7.0f}x" ) print("=" * W) @@ -310,35 +312,30 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: print("Read latency (avg of 5 get_state calls)") print("=" * W) print( - f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'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 in rows: print( f"{turns:>6} {_approx_tokens(turns):>10} " - f"{b_rt * 1000:>10.1f}ms {d_rt * 1000:>10.1f}ms {s_rt * 1000:>10.1f}ms" + f"{b_rt * 1000:>10.1f}ms {d_rt * 1000:>10.1f}ms" ) print("=" * W) print() if storage_results: best = storage_results[-1] - turns, b_bytes, s_bytes, ratio = best - _, _, _, _, b_rt, _, s_rt = rows[-1] + turns, b_bytes, d_bytes, ratio = best + _, _, _, b_rt, d_rt = rows[-1] print( - f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(s_bytes)} ({ratio:.0f}x less storage); " - f"read {b_rt * 1000:.1f}ms → {s_rt * 1000:.1f}ms" + f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(d_bytes)} ({ratio:.0f}x less storage); " + f"read {b_rt * 1000:.1f}ms → {d_rt * 1000:.1f}ms" ) print() print("Legend:") - print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") - print( - " delta = DeltaChannel(add_messages) — O(N) storage, unbounded chain" - ) - print( - f" delta+snap = DeltaChannel(add_messages, snapshot_every={SNAPSHOT_EVERY}) — O(N) storage, bounded read depth" - ) + print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") + print(" delta = DeltaChannel(add_messages) — O(N) storage, reconstructed from writes") print() @@ -359,15 +356,10 @@ def test_delta_channel_benchmark(capsys: Any) -> None: for turns in [25, 50]: _, _, b_bytes = _run_turns(turns, BinaryState) _, _, d_bytes = _run_turns(turns, DeltaState) - _, _, s_bytes = _run_turns(turns, DeltaSnapshotState) assert d_bytes < b_bytes, ( f"DeltaChannel should use less storage at {turns} turns, " f"got delta={d_bytes} binary={b_bytes}" ) - assert s_bytes < b_bytes, ( - f"DeltaChannel+snapshot should use less storage at {turns} turns, " - f"got snapshot={s_bytes} binary={b_bytes}" - ) # ---------------------------------------------------------------------------