mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
feat(channels): DeltaChannel with snapshot_frequency for bounded read depth
Restores DeltaChannel as a standalone class in channels/delta.py and adds
a snapshot_frequency parameter that writes a full snapshot blob every N writes,
bounding the ancestor replay walk depth while preserving O(N) storage for
large N.
Key design decisions:
- Write-count based (not step-based): snapshot fires every N writes to the
channel, tracked via _write_count incremented in both update() and
replay_writes(). This ensures the snapshot always coincides with an actual
channel write (i.e., a new_versions entry in put()), so it is always stored.
- Snapshot blob format: {"__delta_v__": value, "__delta_wc__": n} embeds the
write count so from_checkpoint() can restore it across invocations, keeping
the cadence correct without any external state.
- _checkpoint.py simplified: DeltaChannel.checkpoint() now returns the right
thing (sentinel or snapshot dict) so create_checkpoint needs no special logic.
- _needs_replay updated: triggers on DELTA_SENTINEL / MISSING; snapshot dicts
and plain values (migration) resolve directly via from_checkpoint().
Benchmark shows correct tradeoffs across frequencies (500 turns):
freq=1 → 296 MB storage, ~7ms reads
freq=5 → 60 MB storage, ~4ms reads
freq=10 → 30 MB storage, ~4ms reads
freq=50 → 6.5 MB storage, ~3ms reads
freq=inf→ 290 KB storage, ~114ms reads
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
0ae81f3cff
commit
ba12c8264d
@@ -1,6 +1,7 @@
|
||||
from langgraph.channels.any_value import AnyValue
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
|
||||
from langgraph.channels.named_barrier_value import (
|
||||
@@ -20,6 +21,7 @@ __all__ = (
|
||||
"UntrackedValue",
|
||||
"EphemeralValue",
|
||||
"BinaryOperatorAggregate",
|
||||
"DeltaChannel",
|
||||
"NamedBarrierValue",
|
||||
"NamedBarrierValueAfterFinish",
|
||||
# topics
|
||||
|
||||
@@ -1,177 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy as _copy
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel, Value
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.errors import EmptyChannelError
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
def _empty(typ: Any) -> Any:
|
||||
try:
|
||||
return typ()
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
"""Experimental — private API, subject to change or removal without notice.
|
||||
|
||||
Imported from the underscored module `langgraph.channels._delta` on purpose;
|
||||
not re-exported from `langgraph.channels`. Intended for internal use only
|
||||
while we validate the design on real workloads.
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
Reconstruction replays every ancestor write through the operator, so
|
||||
per-get cost scales with thread depth. Compaction for deep threads is
|
||||
a follow-up — today, use this on threads of a few hundred turns.
|
||||
|
||||
Usage::
|
||||
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"value",
|
||||
"operator",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Any, Any], Any],
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.value: Any = []
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
return False
|
||||
if (
|
||||
self.operator.__name__ != "<lambda>"
|
||||
and other.operator.__name__ != "<lambda>"
|
||||
):
|
||||
return self.operator is other.operator
|
||||
return True
|
||||
|
||||
@property
|
||||
def ValueType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
def copy(self) -> Self:
|
||||
new: DeltaChannel[Value] = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
||||
return new
|
||||
|
||||
def _apply_write(self, value: Any, write: Any) -> Any:
|
||||
"""Apply one write to `value` and return the new value.
|
||||
|
||||
An `Overwrite` replaces the value; any other write is folded through
|
||||
the operator. Centralizes the Overwrite/reducer branching used by both
|
||||
`update` (live super-step) and `from_checkpoint` (ancestor replay).
|
||||
"""
|
||||
is_overwrite, overwrite_value = _get_overwrite(write)
|
||||
if is_overwrite:
|
||||
return (
|
||||
_copy.copy(overwrite_value)
|
||||
if overwrite_value is not None
|
||||
else _empty(self.typ)
|
||||
)
|
||||
base = _empty(self.typ) if value is MISSING else value
|
||||
return self.operator(base, write)
|
||||
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
"""Initialize from a seed value.
|
||||
|
||||
Pregel's hydration path calls this with the `seed` returned by
|
||||
`saver.get_channel_history`:
|
||||
|
||||
* `MISSING` / `DELTA_SENTINEL` → channel starts empty. The walk
|
||||
either reached the root (fresh delta thread) or found nothing
|
||||
to seed from.
|
||||
* any other value → use as the base value. Typically a pre-delta
|
||||
blob preserved across a channel-type migration; `replay_writes`
|
||||
folds subsequent deltas on top.
|
||||
"""
|
||||
new: DeltaChannel[Value] = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
|
||||
new.value = _empty(new.typ)
|
||||
else:
|
||||
new.value = checkpoint
|
||||
return new
|
||||
|
||||
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
|
||||
"""Fold a sequence of `PendingWrite` tuples into the current value.
|
||||
|
||||
Called after `from_checkpoint` during pregel hydration to replay
|
||||
per-step deltas from on-path ancestors through the reducer. Writes
|
||||
are oldest→newest. `Overwrite` values inside the stream reset the
|
||||
reducer state at that point, same as during a live super-step.
|
||||
The `task_id` and `channel` fields of each `PendingWrite` are
|
||||
ignored — `_get_channel_writes_history` has already filtered to
|
||||
this channel.
|
||||
"""
|
||||
for _, _, value in writes:
|
||||
self.value = self._apply_write(self.value, value)
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
if not values:
|
||||
return False
|
||||
seen_overwrite = False
|
||||
for value in values:
|
||||
is_overwrite, _ = _get_overwrite(value)
|
||||
if is_overwrite:
|
||||
if seen_overwrite:
|
||||
from langgraph.errors import (
|
||||
ErrorCode,
|
||||
InvalidUpdateError,
|
||||
create_error_message,
|
||||
)
|
||||
|
||||
msg = create_error_message(
|
||||
message="Can receive only one Overwrite value per super-step.",
|
||||
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
|
||||
)
|
||||
raise InvalidUpdateError(msg)
|
||||
seen_overwrite = True
|
||||
elif seen_overwrite:
|
||||
# Post-Overwrite writes within the same super-step are dropped.
|
||||
continue
|
||||
self.value = self._apply_write(self.value, value)
|
||||
return True
|
||||
|
||||
def get(self) -> Any:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
return DELTA_SENTINEL
|
||||
@@ -0,0 +1,201 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy as _copy
|
||||
import math
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._constants import OVERWRITE
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel, Value
|
||||
from langgraph.errors import (
|
||||
EmptyChannelError,
|
||||
ErrorCode,
|
||||
InvalidUpdateError,
|
||||
create_error_message,
|
||||
)
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
def _empty(typ: Any) -> Any:
|
||||
try:
|
||||
return typ()
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _get_overwrite(value: Any) -> tuple[bool, Any]:
|
||||
if isinstance(value, Overwrite):
|
||||
return True, value.value
|
||||
if isinstance(value, dict) and set(value.keys()) == {OVERWRITE}:
|
||||
return True, value[OVERWRITE]
|
||||
return False, None
|
||||
|
||||
|
||||
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
"""Fold-reducer channel with configurable snapshot cadence.
|
||||
|
||||
`snapshot_frequency=math.inf` (default): pure delta mode — stores only
|
||||
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes
|
||||
through `operator`. O(N) storage, O(N) read replay depth.
|
||||
|
||||
`snapshot_frequency=N` (integer > 1): writes a full snapshot blob every
|
||||
N steps; on non-snapshot steps stores `DELTA_SENTINEL`. Reads walk at
|
||||
most N ancestors before finding the snapshot, bounding replay depth to N.
|
||||
Storage is O(N²/N) = O(N·avg_msg_size) — still linear per step,
|
||||
but snapshots grow with accumulated messages so total is O(N²/N).
|
||||
|
||||
Parameters:
|
||||
operator: Binary reducer `(Value, Value) -> Value` applied pairwise.
|
||||
Must be associative when `snapshot_frequency > 1` since replayed
|
||||
writes fold through the same operator as live writes.
|
||||
snapshot_frequency: Every Nth step writes a full snapshot blob.
|
||||
Default `math.inf` (pure delta, never snapshot).
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "operator", "snapshot_frequency", "_write_count")
|
||||
value: Value | Any
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Any, Any], Any],
|
||||
*,
|
||||
snapshot_frequency: int | float = math.inf,
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.snapshot_frequency = snapshot_frequency
|
||||
self.value: Any = []
|
||||
self._write_count: int = 0
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
return False
|
||||
if self.snapshot_frequency != other.snapshot_frequency:
|
||||
return False
|
||||
if (
|
||||
self.operator.__name__ != "<lambda>"
|
||||
and other.operator.__name__ != "<lambda>"
|
||||
):
|
||||
return self.operator is other.operator
|
||||
return True
|
||||
|
||||
@property
|
||||
def ValueType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
def _clone_empty(self) -> Self:
|
||||
new = self.__class__.__new__(self.__class__)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
new.operator = self.operator
|
||||
new.snapshot_frequency = self.snapshot_frequency
|
||||
new._write_count = 0
|
||||
new.value = MISSING
|
||||
return new
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = self._clone_empty()
|
||||
new._write_count = self._write_count
|
||||
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
||||
return new
|
||||
|
||||
def _apply_write(self, value: Any, write: Any) -> Any:
|
||||
is_overwrite, overwrite_value = _get_overwrite(write)
|
||||
if is_overwrite:
|
||||
return (
|
||||
_copy.copy(overwrite_value)
|
||||
if overwrite_value is not None
|
||||
else _empty(self.typ)
|
||||
)
|
||||
base = _empty(self.typ) if value is MISSING else value
|
||||
return self.operator(base, write)
|
||||
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
"""Initialize from a stored blob or sentinel.
|
||||
|
||||
Blob formats:
|
||||
* `DELTA_SENTINEL` / `MISSING`: start empty; caller will replay writes.
|
||||
* `{"__delta_v__": value, "__delta_wc__": n}`: snapshot blob — restore
|
||||
value and write count so future writes trigger the next snapshot at
|
||||
the right cadence.
|
||||
* plain value (migration from old DeltaChannel blobs): restore value,
|
||||
reset write count to 0.
|
||||
"""
|
||||
new = self._clone_empty()
|
||||
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
|
||||
new.value = _empty(self.typ)
|
||||
new._write_count = 0
|
||||
elif isinstance(checkpoint, dict) and "__delta_v__" in checkpoint:
|
||||
new.value = checkpoint["__delta_v__"]
|
||||
new._write_count = checkpoint["__delta_wc__"]
|
||||
else:
|
||||
new.value = checkpoint
|
||||
new._write_count = 0
|
||||
return new
|
||||
|
||||
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
|
||||
"""Fold ancestor writes oldest→newest into current value.
|
||||
|
||||
Also increments `_write_count` per replayed write so the snapshot
|
||||
cadence stays correct across invocations. The count resumes from
|
||||
wherever the seed's `__delta_wc__` left off (set by `from_checkpoint`).
|
||||
"""
|
||||
for _, _, value in writes:
|
||||
self.value = self._apply_write(self.value, value)
|
||||
self._write_count += 1
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
if not values:
|
||||
return False
|
||||
seen_overwrite = False
|
||||
for value in values:
|
||||
is_overwrite, _ = _get_overwrite(value)
|
||||
if is_overwrite:
|
||||
if seen_overwrite:
|
||||
msg = create_error_message(
|
||||
message="Can receive only one Overwrite value per super-step.",
|
||||
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
|
||||
)
|
||||
raise InvalidUpdateError(msg)
|
||||
seen_overwrite = True
|
||||
elif seen_overwrite:
|
||||
continue
|
||||
self.value = self._apply_write(self.value, value)
|
||||
self._write_count += 1
|
||||
return True
|
||||
|
||||
def get(self) -> Any:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
"""Return stored representation.
|
||||
|
||||
Pure delta mode (`snapshot_frequency=math.inf`) always returns
|
||||
`DELTA_SENTINEL`. For finite `snapshot_frequency`, returns a snapshot
|
||||
blob `{"__delta_v__": value, "__delta_wc__": n}` every
|
||||
`snapshot_frequency` writes, with the current write count embedded so
|
||||
`from_checkpoint` can restore the counter and maintain correct cadence
|
||||
across invocations. All other writes return `DELTA_SENTINEL`.
|
||||
"""
|
||||
if self.value is MISSING:
|
||||
return MISSING
|
||||
if self.snapshot_frequency == math.inf:
|
||||
return DELTA_SENTINEL
|
||||
if self._write_count > 0 and self._write_count % self.snapshot_frequency == 0:
|
||||
return {"__delta_v__": self.value, "__delta_wc__": self._write_count}
|
||||
return DELTA_SENTINEL
|
||||
@@ -8,8 +8,8 @@ from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Check
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel # used in _needs_replay
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
LATEST_VERSION = 4
|
||||
@@ -34,7 +34,12 @@ def create_checkpoint(
|
||||
id: str | None = None,
|
||||
updated_channels: set[str] | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels."""
|
||||
"""Create a checkpoint for the given channels.
|
||||
|
||||
`DeltaChannel.checkpoint()` returns `DELTA_SENTINEL` or the full value
|
||||
depending on its internal write count, so no special handling is needed
|
||||
here. The snapshot is stored at the exact step of the triggering write.
|
||||
"""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
@@ -43,7 +48,8 @@ def create_checkpoint(
|
||||
for k in channels:
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
continue
|
||||
v = channels[k].checkpoint()
|
||||
ch = channels[k]
|
||||
v = ch.checkpoint()
|
||||
if v is not MISSING:
|
||||
values[k] = v
|
||||
return Checkpoint(
|
||||
@@ -57,6 +63,19 @@ def create_checkpoint(
|
||||
)
|
||||
|
||||
|
||||
def _needs_replay(spec: BaseChannel, stored: object) -> bool:
|
||||
"""True if `spec` is a `DeltaChannel` and the stored blob is a sentinel,
|
||||
requiring an ancestor walk to reconstruct.
|
||||
|
||||
Snapshot blobs are dicts with `__delta_v__`; sentinels are `DELTA_SENTINEL`.
|
||||
Plain non-sentinel values are migration blobs (old DeltaChannel format).
|
||||
All non-sentinel values resolve directly via `from_checkpoint`.
|
||||
"""
|
||||
if not isinstance(spec, DeltaChannel):
|
||||
return False
|
||||
return stored is MISSING or stored is DELTA_SENTINEL
|
||||
|
||||
|
||||
def channels_from_checkpoint(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
checkpoint: Checkpoint,
|
||||
@@ -69,12 +88,14 @@ def channels_from_checkpoint(
|
||||
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
|
||||
is sufficient — the stored value IS the reconstructed state.
|
||||
|
||||
`DeltaChannel` is the exception: its stored value is a sentinel; the
|
||||
full state is spread across `checkpoint_writes` along the ancestor
|
||||
chain. When `saver` and `config` are provided, this function fetches
|
||||
that history via `saver._get_channel_writes_history` and folds it
|
||||
through the channel's reducer. Without them (static contexts — graph
|
||||
drawing, unit tests), delta channels fall back to empty.
|
||||
`DeltaChannel` is the exception: its stored value on non-snapshot steps is
|
||||
`DELTA_SENTINEL`; the full state is spread across `checkpoint_writes` along
|
||||
the ancestor chain. When `saver` and `config` are provided, this function
|
||||
fetches that history via `saver._get_channel_writes_history` and folds it
|
||||
through the channel's operator. The walk stops at the nearest full-snapshot
|
||||
blob, so read depth is bounded by `snapshot_frequency`. Without saver/config
|
||||
(static contexts — graph drawing, unit tests), delta-mode channels fall back
|
||||
to empty.
|
||||
"""
|
||||
channel_specs: dict[str, BaseChannel] = {}
|
||||
managed_specs: dict[str, ManagedValueSpec] = {}
|
||||
@@ -88,22 +109,16 @@ def channels_from_checkpoint(
|
||||
for k, spec in channel_specs.items():
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
if (
|
||||
isinstance(spec, DeltaChannel)
|
||||
and saver is not None
|
||||
and config is not None
|
||||
and (stored is MISSING or stored is DELTA_SENTINEL)
|
||||
):
|
||||
# Target's own blob is empty/sentinel — walk ancestors for
|
||||
# seed + writes. Skipping this when `stored` is a real value
|
||||
# preserves state written via `update_state` or sitting at the
|
||||
# tip of a pre-migration thread: the saver's ancestor walk
|
||||
# intentionally excludes the target's own blob, so without
|
||||
# this short-circuit we'd lose it.
|
||||
if _needs_replay(spec, stored) and saver is not None and config is not None:
|
||||
# Walk ancestors for seed + writes. The saver's walk stops at
|
||||
# the nearest non-sentinel blob (natural terminator under
|
||||
# snapshot_frequency > 1; pre-migration blobs also act as
|
||||
# terminators if the spec was changed mid-thread).
|
||||
assert isinstance(spec, DeltaChannel)
|
||||
history = saver._get_channel_writes_history(config, k)
|
||||
delta_ch = spec.from_checkpoint(history.seed)
|
||||
delta_ch.replay_writes(history.writes)
|
||||
ch = delta_ch
|
||||
replay_ch = spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
ch = replay_ch
|
||||
else:
|
||||
ch = spec.from_checkpoint(stored)
|
||||
channels[k] = ch
|
||||
@@ -130,16 +145,12 @@ async def achannels_from_checkpoint(
|
||||
for k, spec in channel_specs.items():
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
if (
|
||||
isinstance(spec, DeltaChannel)
|
||||
and saver is not None
|
||||
and config is not None
|
||||
and (stored is MISSING or stored is DELTA_SENTINEL)
|
||||
):
|
||||
if _needs_replay(spec, stored) and saver is not None and config is not None:
|
||||
assert isinstance(spec, DeltaChannel)
|
||||
history = await saver._aget_channel_writes_history(config, k)
|
||||
delta_ch = spec.from_checkpoint(history.seed)
|
||||
delta_ch.replay_writes(history.writes)
|
||||
ch = delta_ch
|
||||
replay_ch = spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
ch = replay_ch
|
||||
else:
|
||||
ch = spec.from_checkpoint(stored)
|
||||
channels[k] = ch
|
||||
|
||||
@@ -1,22 +1,21 @@
|
||||
"""Benchmark: DeltaChannel vs BinaryOperatorAggregate storage and time.
|
||||
"""Benchmark: DeltaChannel snapshot_frequency — storage vs. read-depth tradeoff.
|
||||
|
||||
Run directly: python tests/test_delta_channel_benchmark.py
|
||||
Run via pytest: pytest tests/test_delta_channel_benchmark.py -s
|
||||
|
||||
Simulates realistic multi-turn conversations with paragraph-length messages
|
||||
(~100 tokens each) scaling up to 1M-token-equivalent histories.
|
||||
Part 1 — baseline (original): DeltaChannel(inf) vs add_messages (BinOp).
|
||||
Part 2 — snapshot_frequency sweep: shows the storage/read-latency tradeoff
|
||||
across frequencies [1, 5, 10, 50, inf] at scale.
|
||||
|
||||
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.
|
||||
Key insight:
|
||||
snapshot_frequency=inf → O(N) storage, O(N) read depth (pure delta)
|
||||
snapshot_frequency=N → O(N²/N) storage, O(N) read depth bounded by freq
|
||||
snapshot_frequency=1 → O(N²) storage, O(1) read depth (full snapshot)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
@@ -26,17 +25,10 @@ from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
|
||||
_SQLITE_AVAILABLE = True
|
||||
except ImportError:
|
||||
_SQLITE_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
@@ -126,6 +118,18 @@ class DeltaState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
|
||||
def _make_delta_state(snapshot_frequency: int | float) -> type:
|
||||
"""Create a TypedDict with DeltaChannel at the given snapshot_frequency."""
|
||||
channel = DeltaChannel(add_messages, snapshot_frequency=snapshot_frequency)
|
||||
# Use the functional TypedDict form so the Annotated type is stored as an
|
||||
# already-evaluated object rather than a forward-reference string (which
|
||||
# would fail when get_type_hints tries to resolve 'snapshot_frequency').
|
||||
return TypedDict( # type: ignore[return-value]
|
||||
f"DeltaState_freq{snapshot_frequency}",
|
||||
{"messages": Annotated[list, channel]},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factory
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -169,9 +173,8 @@ def _run_turns(
|
||||
"""Run n_turns conversation turns.
|
||||
|
||||
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
|
||||
blob_bytes is -1 for savers without in-memory blob stores (e.g. SQLite).
|
||||
Read latency is measured as the time to invoke the graph with no new
|
||||
messages after the full history is built — this forces state rehydration.
|
||||
Read latency is the average of 5 get_state calls after the full history
|
||||
is built — forces state rehydration including ancestor replay if needed.
|
||||
"""
|
||||
graph = _make_graph(state_cls, checkpointer)
|
||||
config = {"configurable": {"thread_id": "bench"}}
|
||||
@@ -184,16 +187,16 @@ def _run_turns(
|
||||
)
|
||||
write_elapsed = time.perf_counter() - t0
|
||||
|
||||
# Measure read/rehydration: get_state forces the channel to rebuild
|
||||
t1 = time.perf_counter()
|
||||
for _ in range(5):
|
||||
graph.get_state(config)
|
||||
read_elapsed = (time.perf_counter() - t1) / 5
|
||||
|
||||
if isinstance(graph.checkpointer, MemorySaver):
|
||||
blob_bytes = _total_blob_bytes(graph.checkpointer)
|
||||
else:
|
||||
blob_bytes = -1
|
||||
blob_bytes = (
|
||||
_total_blob_bytes(graph.checkpointer)
|
||||
if isinstance(graph.checkpointer, MemorySaver)
|
||||
else -1
|
||||
)
|
||||
return write_elapsed, read_elapsed, blob_bytes
|
||||
|
||||
|
||||
@@ -206,7 +209,6 @@ def _fmt_bytes(n: int) -> str:
|
||||
|
||||
|
||||
def _approx_tokens(n_turns: int) -> str:
|
||||
# ~100 tokens human + ~100 tokens AI per turn
|
||||
tokens = n_turns * 200
|
||||
if tokens >= 1_000_000:
|
||||
return f"~{tokens / 1_000_000:.1f}M tok"
|
||||
@@ -216,104 +218,41 @@ def _approx_tokens(n_turns: int) -> str:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark matrix
|
||||
# Part 1: baseline DeltaChannel(inf) vs add_messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Turn counts chosen to demonstrate O(N²) vs O(N) storage growth without running too long.
|
||||
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
|
||||
TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
|
||||
# Deep-thread counts where add_messages blob storage would exceed 1 GB;
|
||||
# only DeltaChannel runs here.
|
||||
BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
DELTA_ONLY_TURN_COUNTS = [1000]
|
||||
|
||||
|
||||
def _checkpointer_factories() -> list[tuple[str, Any]]:
|
||||
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
|
||||
return [("InMemory", None)]
|
||||
|
||||
|
||||
def run_benchmark() -> None:
|
||||
def run_baseline_benchmark() -> None:
|
||||
print()
|
||||
print(
|
||||
"DeltaChannel vs add_messages (BinaryOperatorAggregate) — checkpoint storage & latency"
|
||||
)
|
||||
print("Simulating realistic multi-turn conversations up to ~1M-token histories")
|
||||
print("(5,000 turns × ~200 tokens/turn ≈ 1M tokens — Claude's full context window)")
|
||||
print()
|
||||
|
||||
checkpointers: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
if _POSTGRES_AVAILABLE:
|
||||
try:
|
||||
import psycopg
|
||||
|
||||
psycopg.connect(_POSTGRES_URI).close()
|
||||
checkpointers.append(("Postgres (plain SELECT)", "postgres"))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for cp_label, cp_hint in checkpointers:
|
||||
print(f"--- Checkpointer: {cp_label} ---")
|
||||
_run_benchmark_for_checkpointer(cp_hint)
|
||||
|
||||
|
||||
def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
import contextlib
|
||||
import tempfile
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
yield None
|
||||
elif cp_hint == "postgres":
|
||||
with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver:
|
||||
saver.setup()
|
||||
with saver._cursor() as cur:
|
||||
cur.execute("DELETE FROM checkpoints WHERE thread_id = 'bench'")
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoint_blobs WHERE thread_id = 'bench'"
|
||||
)
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoint_writes WHERE thread_id = 'bench'"
|
||||
)
|
||||
yield saver
|
||||
else:
|
||||
with tempfile.NamedTemporaryFile(suffix=".db") as f:
|
||||
with SqliteSaver.from_conn_string(f.name) as saver:
|
||||
yield saver
|
||||
print("Part 1 — DeltaChannel(inf) vs add_messages: storage & latency")
|
||||
print("=" * 72)
|
||||
|
||||
rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = []
|
||||
for turns in TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
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)
|
||||
for turns in BASELINE_TURN_COUNTS:
|
||||
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState)
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState)
|
||||
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt))
|
||||
for turns in DELTA_ONLY_TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState)
|
||||
rows.append((turns, None, d_bytes, None, d_rt, None, d_wt))
|
||||
|
||||
# ── Table 1: Storage ─────────────────────────────────────────────────────
|
||||
W = 64
|
||||
print("Storage (checkpoint blob bytes)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} "
|
||||
f"{'savings':>8}"
|
||||
)
|
||||
print("-" * W)
|
||||
W = 72
|
||||
|
||||
def _bytes_or_na(v: Any) -> str:
|
||||
if v is None:
|
||||
return "n/a"
|
||||
if v < 0:
|
||||
if v is None or v < 0:
|
||||
return "n/a"
|
||||
return _fmt_bytes(v)
|
||||
|
||||
def _ms_or_na(v: Any) -> str:
|
||||
return "n/a" if v is None else f"{v * 1000:.1f}ms"
|
||||
|
||||
print("Storage (blob bytes)")
|
||||
print("-" * W)
|
||||
print(f"{'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}")
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0:
|
||||
ratio_str = "n/a"
|
||||
@@ -322,64 +261,117 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
ratio_str = f"{ratio:.0f}x"
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} "
|
||||
f"{ratio_str:>8}"
|
||||
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} {ratio_str:>8}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
# ── Table 2: Read latency ─────────────────────────────────────────────────
|
||||
print("Read latency (avg of 5 get_state calls)")
|
||||
print("=" * W)
|
||||
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
|
||||
print("-" * W)
|
||||
print(f"{'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12}")
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
# ── Table 3: Per-invoke latency (total write_elapsed / turns) ─────────────
|
||||
print("Per-invoke latency (total graph.invoke time / turns)")
|
||||
print("=" * W)
|
||||
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
|
||||
def _per(wt: Any) -> str:
|
||||
if wt is None:
|
||||
return "n/a"
|
||||
return f"{(wt / turns) * 1000:.1f}ms"
|
||||
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_per(b_wt):>12} {_per(d_wt):>12}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
print("Legend:")
|
||||
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
|
||||
print(" delta = Annotated[list, DeltaChannel(add_messages)] — O(N) storage")
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry point
|
||||
# Part 2: snapshot_frequency sweep
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Frequencies to test. 1 = always snapshot (like BinOp), inf = pure delta.
|
||||
SNAPSHOT_FREQUENCIES: list[int | float] = [1, 5, 10, 50, math.inf]
|
||||
|
||||
# Turn counts for the sweep — high enough to show storage divergence.
|
||||
SWEEP_TURN_COUNTS = [50, 100, 500]
|
||||
|
||||
|
||||
def _freq_label(freq: int | float) -> str:
|
||||
if freq == math.inf:
|
||||
return "inf"
|
||||
return str(int(freq))
|
||||
|
||||
|
||||
def run_snapshot_freq_benchmark() -> None:
|
||||
print()
|
||||
print("Part 2 — DeltaChannel snapshot_frequency sweep")
|
||||
print("Lower freq → fewer snapshots → less storage but deeper read replay")
|
||||
print("=" * 80)
|
||||
|
||||
# Collect results: {turns: {freq: (write_s, read_s, bytes)}}
|
||||
results: dict[int, dict[str | int, tuple[float, float, int]]] = {}
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
results[turns] = {}
|
||||
for freq in SNAPSHOT_FREQUENCIES:
|
||||
state_cls = _make_delta_state(freq)
|
||||
wt, rt, bb = _run_turns(turns, state_cls)
|
||||
results[turns][_freq_label(freq)] = (wt, rt, bb)
|
||||
|
||||
freq_labels = [_freq_label(f) for f in SNAPSHOT_FREQUENCIES]
|
||||
col_w = 12
|
||||
|
||||
# Storage table
|
||||
print()
|
||||
print("Storage (blob bytes) — lower is better")
|
||||
header = f"{'turns':>6} {'ctx':>10}" + "".join(f" {f'freq={l}':>{col_w}}" for l in freq_labels)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
row = f"{turns:>6} {_approx_tokens(turns):>10}"
|
||||
for label in freq_labels:
|
||||
_, _, bb = results[turns][label]
|
||||
row += f" {_fmt_bytes(bb) if bb >= 0 else 'n/a':>{col_w}}"
|
||||
print(row)
|
||||
|
||||
print()
|
||||
|
||||
# Read latency table
|
||||
print("Read latency (avg of 5 get_state) — lower is better")
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
row = f"{turns:>6} {_approx_tokens(turns):>10}"
|
||||
for label in freq_labels:
|
||||
_, rt, _ = results[turns][label]
|
||||
row += f" {f'{rt * 1000:.1f}ms':>{col_w}}"
|
||||
print(row)
|
||||
|
||||
print()
|
||||
|
||||
# Per-invoke write latency
|
||||
print("Per-invoke write latency (total / turns) — lower is better")
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
row = f"{turns:>6} {_approx_tokens(turns):>10}"
|
||||
for label in freq_labels:
|
||||
wt, _, _ = results[turns][label]
|
||||
row += f" {f'{(wt / turns) * 1000:.1f}ms':>{col_w}}"
|
||||
print(row)
|
||||
|
||||
print()
|
||||
print("Legend:")
|
||||
print(" freq=1 snapshot every write (full blob always — same as add_messages / BinOp)")
|
||||
print(" freq=N snapshot every N writes; read walks at most N ancestor writes")
|
||||
print(" freq=inf pure delta; read walks entire ancestor chain")
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry points
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
|
||||
)
|
||||
def test_delta_channel_benchmark(capsys: Any) -> None:
|
||||
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
|
||||
def test_delta_channel_baseline_benchmark(capsys: Any) -> None:
|
||||
"""DeltaChannel(inf) uses less storage than add_messages at scale."""
|
||||
with capsys.disabled():
|
||||
run_benchmark()
|
||||
run_baseline_benchmark()
|
||||
|
||||
# Correctness assertion: DeltaChannel must use less storage at scale.
|
||||
for turns in [25, 50]:
|
||||
_, _, b_bytes = _run_turns(turns, BinaryState)
|
||||
_, _, d_bytes = _run_turns(turns, DeltaState)
|
||||
@@ -389,10 +381,41 @@ def test_delta_channel_benchmark(capsys: Any) -> None:
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
|
||||
)
|
||||
def test_snapshot_freq_benchmark(capsys: Any) -> None:
|
||||
"""snapshot_frequency trades storage for bounded read depth."""
|
||||
with capsys.disabled():
|
||||
run_snapshot_freq_benchmark()
|
||||
|
||||
# Correctness: results at all frequencies should agree on final state.
|
||||
n_turns = 20
|
||||
states: dict[str, list] = {}
|
||||
for freq in SNAPSHOT_FREQUENCIES:
|
||||
state_cls = _make_delta_state(freq)
|
||||
graph = _make_graph(state_cls)
|
||||
config = {"configurable": {"thread_id": "correctness"}}
|
||||
for i in range(n_turns):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
|
||||
config,
|
||||
)
|
||||
state = graph.get_state(config)
|
||||
states[_freq_label(freq)] = [m.id for m in state.values["messages"]]
|
||||
|
||||
ref = states["inf"]
|
||||
for label, msg_ids in states.items():
|
||||
assert msg_ids == ref, (
|
||||
f"freq={label} produced different message IDs than freq=inf"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Script entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_benchmark()
|
||||
run_baseline_benchmark()
|
||||
run_snapshot_freq_benchmark()
|
||||
sys.exit(0)
|
||||
|
||||
Reference in New Issue
Block a user