feat(delta-channel): store sentinel in blobs, reconstruct from checkpoint_writes

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.
This commit is contained in:
Sydney Runkle
2026-04-30 14:49:05 -04:00
parent 5d72a3f390
commit 4a6ddbb0e3
12 changed files with 410 additions and 603 deletions
@@ -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": {
@@ -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
)
@@ -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,
@@ -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.
@@ -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=(
@@ -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"),
}
)
@@ -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:
+4 -5
View File
@@ -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)
+45 -42
View File
@@ -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"}]
+29 -89
View File
@@ -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__ != "<lambda>"
and other.operator.__name__ != "<lambda>"
@@ -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()
+70 -193
View File
@@ -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"}
@@ -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}"
)
# ---------------------------------------------------------------------------