diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index a69c230b5..cc7fcc471 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -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 diff --git a/libs/langgraph/langgraph/channels/_delta.py b/libs/langgraph/langgraph/channels/_delta.py deleted file mode 100644 index 3239d237b..000000000 --- a/libs/langgraph/langgraph/channels/_delta.py +++ /dev/null @@ -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__ != "" - and other.operator.__name__ != "" - ): - 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 diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py new file mode 100644 index 000000000..992f49a92 --- /dev/null +++ b/libs/langgraph/langgraph/channels/delta.py @@ -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__ != "" + and other.operator.__name__ != "" + ): + 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 diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 51350071d..3f56a09d8 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -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 diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index bc393d804..379b12f7e 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -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)