From ee5b3fb4ca84e9692eaf5b8a5e80020fbe2dffc5 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 24 Apr 2026 11:10:19 -0400 Subject: [PATCH] feat(channels): AggregateChannel unifies BinOp + DeltaChannel with snapshot_frequency knob MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New `AggregateChannel(operator, *, snapshot_frequency=1, typ=None)` replaces the experimental `DeltaChannel`. `snapshot_frequency=1` (default) is today's `BinaryOperatorAggregate`. `snapshot_frequency=math.inf` is today's pure-delta behavior. Integer values between bound replay depth — deep-thread reads become O(snapshot_frequency) instead of O(thread depth). - `BinaryOperatorAggregate` is now a thin subclass (snapshot_frequency=1), preserving isinstance checks and `_is_field_binop` detection. - `DeltaChannel` (private, experimental, underscored module) is removed. Migration: `AggregateChannel(op, snapshot_frequency=math.inf)`. - `create_checkpoint` is step-aware: non-snapshot steps store DELTA_SENTINEL; snapshot steps store the full blob. - `_get_channel_writes_history` walk is fixed: pending_writes of a terminator ancestor encode its state→child transition and are now collected BEFORE checking the blob (old code silently dropped them, which was hidden because no prior scenario had a FULL blob mid-thread). Saver API is unchanged. Batched multi-channel walks, `walk_writes`/ `put_channel_snapshot` refactor, `coalesce=` kwarg, and Option A (channel_versions delta-encoding) are deferred per spec. Design: `docs/superpowers/specs/2026-04-24-aggregate-channel-design.md`. Verified: - snapshot_frequency ∈ {1, 2, 3, 5, 10, math.inf} all reconstruct correctly on a 7-invoke / 14-message thread. - Multi-channel graph with different snapshot_frequency per channel works. - All 10 migration tests pass unchanged (pre-BinOp-to-Delta migration path). - Channel unit tests pass (35 tests). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../2026-04-24-aggregate-channel-design.md | 2 +- .../langgraph/checkpoint/postgres/base.py | 12 +- .../langgraph/checkpoint/base/__init__.py | 32 ++- .../langgraph/checkpoint/memory/__init__.py | 43 ++- libs/langgraph/langgraph/channels/__init__.py | 2 + libs/langgraph/langgraph/channels/_delta.py | 177 ------------ .../langgraph/langgraph/channels/aggregate.py | 256 ++++++++++++++++++ libs/langgraph/langgraph/channels/binop.py | 129 ++------- libs/langgraph/langgraph/graph/state.py | 12 +- .../langgraph/langgraph/pregel/_checkpoint.py | 75 ++--- libs/langgraph/tests/test_channels.py | 18 +- .../tests/test_delta_channel_benchmark.py | 7 +- .../tests/test_delta_channel_migration.py | 7 +- libs/langgraph/tests/test_pregel.py | 12 +- 14 files changed, 405 insertions(+), 379 deletions(-) delete mode 100644 libs/langgraph/langgraph/channels/_delta.py create mode 100644 libs/langgraph/langgraph/channels/aggregate.py diff --git a/docs/superpowers/specs/2026-04-24-aggregate-channel-design.md b/docs/superpowers/specs/2026-04-24-aggregate-channel-design.md index 1d9412f3d..626e2e2ae 100644 --- a/docs/superpowers/specs/2026-04-24-aggregate-channel-design.md +++ b/docs/superpowers/specs/2026-04-24-aggregate-channel-design.md @@ -2,7 +2,7 @@ **Status:** MVP scope approved. Implementation starting 2026-04-24. **Supersedes:** `langgraph.channels._delta.DeltaChannel` (experimental, private). -**Branch:** `delta-channel-writes-based`. +**Branch:** `sr/even-better-writes-idea` (forked from `delta-channel-writes-based`). ## Problem diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index b56c503fd..1b2be664b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -281,7 +281,14 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): collected: list[PendingWrite] = [] # newest first; reversed at the end for cid in ancestors: - # Pre-delta blob terminator: subsumes any writes at this ancestor. + # Collect this ancestor's pending_writes FIRST. They encode the + # transition from state-AT-this-ancestor to state-AT-its-child; + # the ancestor's blob only reflects state AT the ancestor, not + # post-transition. Both pre-delta migration and snapshot-cadence + # cases require these writes folded onto the seed. + for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []): + val = self.serde.loads_typed((type_tag, write_blob)) + collected.append((task_id, channel, val)) ver = ver_of.get(cid) if ver is not None: seed_blob = blob_by_ver.get(ver) @@ -290,9 +297,6 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): if blob_value is not DELTA_SENTINEL: collected.reverse() return _ChannelWritesHistory(seed=blob_value, writes=collected) - for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []): - val = self.serde.loads_typed((type_tag, write_blob)) - collected.append((task_id, channel, val)) collected.reverse() # oldest → newest return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index af4ab3dc4..6f01d6bdc 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -540,13 +540,12 @@ class BaseCheckpointSaver(Generic[V]): tup = self.get_tuple(cursor_config) if tup is None: break - # Pre-delta seed terminator: if the ancestor has a stored - # (non-sentinel) value for this channel, that snapshot - # subsumes any earlier writes on the chain. Stop here. - ancestor_value = tup.checkpoint["channel_values"].get(channel) - if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: - collected.reverse() - return _ChannelWritesHistory(seed=ancestor_value, writes=collected) + # Collect this ancestor's pending_writes FIRST. They encode + # the transition from this ancestor's state to its child's + # state (K → K+1); the ancestor's `channel_values[channel]` + # blob reflects state AT K only, not post-transition. Both + # the pre-delta migration case and the snapshot-cadence + # case require these writes to be included. if tup.pending_writes: # Within a superstep, pending_writes are oldest→newest; # reverse to scan newest-first. @@ -554,6 +553,15 @@ class BaseCheckpointSaver(Generic[V]): if write[1] != channel: continue collected.append(write) + # Seed terminator: any non-sentinel blob on an ancestor + # establishes the reconstruction base (state AT K). The + # writes we just collected fold on top to produce state + # at K+1; subsequent (already-collected, child-side) + # writes fold the chain up to the target. + ancestor_value = tup.checkpoint["channel_values"].get(channel) + if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: + collected.reverse() + return _ChannelWritesHistory(seed=ancestor_value, writes=collected) cursor_config = tup.parent_config collected.reverse() return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) @@ -578,15 +586,17 @@ class BaseCheckpointSaver(Generic[V]): tup = await self.aget_tuple(cursor_config) if tup is None: break - ancestor_value = tup.checkpoint["channel_values"].get(channel) - if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: - collected.reverse() - return _ChannelWritesHistory(seed=ancestor_value, writes=collected) + # See sync variant for rationale: collect pending_writes + # BEFORE checking the blob terminator. if tup.pending_writes: for write in reversed(tup.pending_writes): if write[1] != channel: continue collected.append(write) + ancestor_value = tup.checkpoint["channel_values"].get(channel) + if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: + collected.reverse() + return _ChannelWritesHistory(seed=ancestor_value, writes=collected) cursor_config = tup.parent_config collected.reverse() return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 39bb133dc..cf6901b1c 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -161,19 +161,26 @@ class InMemorySaver( chain.append(current) _, _, parent = entry current = parent - # Scan newest→oldest. A pre-delta blob on an ancestor terminates the - # walk and is bound as `seed`; without this, a thread migrated from - # pre-delta storage would replay ancestor writes all the way to the - # root AND miss any value that lived only in the old blob (e.g. from - # `update_state`). - # - # At each ancestor, check the blob BEFORE processing its pending - # writes: a pre-delta blob represents the state AT that ancestor, - # which already subsumes any writes stored under it. Processing - # those writes first would fold them into the reconstructed value - # twice (once via the blob, once via replay). + # Scan newest→oldest. At each ancestor, collect its pending_writes + # BEFORE checking for a non-sentinel blob terminator. Rationale: + # a blob at ancestor K represents state AT step K; that ancestor's + # pending_writes are the writes that transition state K → state K+1. + # Both the pre-delta migration case (pre-migration blob + post-delta + # child) and the snapshot-cadence case (FULL blob mid-thread + later + # sentinel checkpoints) need those writes folded onto the seed. collected: list[PendingWrite] = [] # newest first for cp_id in chain: # newest → oldest + step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) + # Within a superstep, sorted by (task_id, idx) = oldest → newest; + # reverse for newest-first scan. + for (_task_id, _idx), (tid, ch, serialized, _) in sorted( + step_writes.items(), reverse=True + ): + if ch != channel: + continue + val = self.serde.loads_typed(serialized) + collected.append((tid, ch, val)) + entry = ns_storage.get(cp_id) if entry is not None: ckpt = self.serde.loads_typed(entry[0]) @@ -185,23 +192,13 @@ class InMemorySaver( if blob_entry is not None and blob_entry[0] != "empty": blob_value = self.serde.loads_typed(blob_entry) if blob_value is not DELTA_SENTINEL: - # Pre-delta snapshot terminator. Skip this - # ancestor's writes — the blob subsumes them. + # Non-sentinel blob terminator: state AT this + # ancestor becomes the reconstruction seed. collected.reverse() return _ChannelWritesHistory( seed=blob_value, writes=collected ) - step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) - # Within a superstep, sorted by (task_id, idx) = oldest → newest; - # reverse for newest-first scan. - for (_task_id, _idx), (tid, ch, serialized, _) in sorted( - step_writes.items(), reverse=True - ): - if ch != channel: - continue - val = self.serde.loads_typed(serialized) - collected.append((tid, ch, val)) collected.reverse() return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index a69c230b5..48fb031ce 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,3 +1,4 @@ +from langgraph.channels.aggregate import AggregateChannel from langgraph.channels.any_value import AnyValue from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -19,6 +20,7 @@ __all__ = ( "LastValueAfterFinish", "UntrackedValue", "EphemeralValue", + "AggregateChannel", "BinaryOperatorAggregate", "NamedBarrierValue", "NamedBarrierValueAfterFinish", 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/aggregate.py b/libs/langgraph/langgraph/channels/aggregate.py new file mode 100644 index 000000000..2a22a615b --- /dev/null +++ b/libs/langgraph/langgraph/channels/aggregate.py @@ -0,0 +1,256 @@ +from __future__ import annotations + +import collections.abc +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 NotRequired, Required, 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__ = ("AggregateChannel",) + + +def _strip_extras(t: Any) -> Any: + """Strips Annotated, Required, and NotRequired wrappers.""" + if hasattr(t, "__origin__"): + return _strip_extras(t.__origin__) + if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired): + return _strip_extras(t.__args__[0]) + return t + + +def _concrete(typ: Any) -> Any: + """Replace abstract collection types from `typing`/`collections.abc` with + their instantiable counterparts.""" + typ = _strip_extras(typ) + if typ in (collections.abc.Sequence, collections.abc.MutableSequence): + return list + if typ in (collections.abc.Set, collections.abc.MutableSet): + return set + if typ in (collections.abc.Mapping, collections.abc.MutableMapping): + return dict + return typ + + +def _empty(typ: Any) -> Any: + try: + return typ() + except Exception: + return [] + + +def _get_overwrite(value: Any) -> tuple[bool, Any]: + """Return (is_overwrite, overwrite_value) for an incoming write.""" + 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 AggregateChannel(Generic[Value], BaseChannel[Value, Value, Any]): + """Fold-reducer channel with configurable snapshot cadence. + + `snapshot_frequency=1` (default) writes a full blob every step — same + storage behavior as the classic `BinaryOperatorAggregate`. + + `snapshot_frequency=N` (integer > 1) writes a sentinel on non-snapshot + steps; the value is reconstructed at read time by folding ancestor + writes through `operator`. On every Nth step a full blob is written, + bounding replay depth to N. + + `snapshot_frequency=math.inf` never writes a blob — pure delta storage. + Reconstruction replays every write from thread start. + + Parameters: + operator: Binary reducer `(Value, Value) -> Value` applied pairwise + to accumulate writes. Must be associative for correctness under + `snapshot_frequency != 1` where the fold order across ancestor + replay vs live writes differs from the classic single-fold path. + Most practical reducers (`operator.add`, `add_messages`) satisfy + this. + snapshot_frequency: Every Nth step writes a full snapshot blob. + Default 1 (snapshot always). `math.inf` for pure-delta mode. + Reading at step M with `snapshot_frequency=N` walks at most + `M % N` ancestor writes — bounded replay regardless of thread + depth. + typ: Value type. When used as an `Annotated[T, AggregateChannel(...)]` + state field, the type is inferred from `T` and this kwarg is + unused. Explicit kwarg is the escape hatch for imperative + graph construction. + + Experimental under `snapshot_frequency > 1`: the sentinel+replay path + is the same mechanism that the (now removed) `DeltaChannel` used; the + cadence knob is new and should be validated on real workloads before + being relied on in production. + """ + + __slots__ = ("value", "operator", "snapshot_frequency", "_typ_provided") + + def __init__( + self, + operator: Callable[[Value, Value], Value], + *, + snapshot_frequency: int | float = 1, + typ: type[Value] | None = None, + ) -> None: + self._typ_provided = typ is not None + concrete_typ = _concrete(typ) if typ is not None else list + super().__init__(concrete_typ) + self.operator = operator + self.snapshot_frequency = snapshot_frequency + try: + self.value = concrete_typ() + except Exception: + self.value = MISSING + + def __eq__(self, other: object) -> bool: + if not isinstance(other, AggregateChannel): + 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: + """Create a blank clone preserving all attributes, bypassing __init__. + + Subclasses (e.g. BinaryOperatorAggregate) have different __init__ + signatures; going through __init__ from copy/from_checkpoint would + pass kwargs those subclasses don't accept. Bypassing avoids that. + """ + 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._typ_provided = self._typ_provided + new.value = MISSING + return new + + def copy(self) -> Self: + new = self._clone_empty() + new.value = self.value if self.value is MISSING else _copy.copy(self.value) + return new + + def is_snapshot_step(self, step: int) -> bool: + """Return True if a full blob should be written at this step. + + `snapshot_frequency=1` → always. `snapshot_frequency=math.inf` → + never. Otherwise, `step % snapshot_frequency == 0`. + """ + if self.snapshot_frequency == 1: + return True + if self.snapshot_frequency == math.inf: + return False + return step % self.snapshot_frequency == 0 + + def _apply_write(self, value: Any, write: Any) -> Any: + """Apply one write and return the new value. Handles Overwrite.""" + 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) + ) + if value is MISSING: + return write + return self.operator(value, write) + + def from_checkpoint(self, checkpoint: Any) -> Self: + """Initialize from a stored blob, sentinel, or MISSING. + + If the stored value is a full blob, use it as-is. If it is + `DELTA_SENTINEL` or `MISSING`, start empty — the caller (pregel) + is responsible for replaying writes via `replay_writes` when + applicable. + """ + new = self._clone_empty() + if checkpoint is MISSING or checkpoint is DELTA_SENTINEL: + new.value = _empty(self.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 by pregel after `from_checkpoint(seed)` to replay per-step + deltas from on-path ancestors through the operator. Writes are + oldest→newest. Overwrite markers reset the reducer state at that + point. `task_id` and `channel` fields are ignored — the caller + has already filtered to this channel. + """ + for _, _, value in writes: + self.value = self._apply_write(self.value, value) + + def update(self, values: Sequence[Value]) -> 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) + return True + + def get(self) -> Value: + 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 the serializable representation of current state. + + For `snapshot_frequency=math.inf` (pure delta), always returns + `DELTA_SENTINEL` — the value lives in `checkpoint_writes` and is + reconstructed by replay, never materialized as a blob. + + For integer `snapshot_frequency`, returns the full value. Pregel's + `create_checkpoint` is responsible for consulting + `is_snapshot_step(step)` to decide whether to actually store the + full value or write `DELTA_SENTINEL` for that step. + """ + if self.value is MISSING: + return MISSING + if self.snapshot_frequency == math.inf: + return DELTA_SENTINEL + return self.value diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index b33ea52fa..ab07c6ae7 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/libs/langgraph/langgraph/channels/binop.py @@ -1,44 +1,21 @@ -import collections.abc -from collections.abc import Callable, Sequence -from typing import Any, Generic +from collections.abc import Callable +from typing import Generic -from typing_extensions import NotRequired, Required, 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.channels.aggregate import ( + AggregateChannel, ) -from langgraph.types import Overwrite +from langgraph.channels.aggregate import ( + _get_overwrite as _get_overwrite, +) +from langgraph.channels.aggregate import ( + _strip_extras as _strip_extras, +) +from langgraph.channels.base import Value __all__ = ("BinaryOperatorAggregate",) -# Adapted from typing_extensions -def _strip_extras(t): # type: ignore[no-untyped-def] - """Strips Annotated, Required and NotRequired from a given type.""" - if hasattr(t, "__origin__"): - return _strip_extras(t.__origin__) - if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired): - return _strip_extras(t.__args__[0]) - - return t - - -def _get_overwrite(value: Any) -> tuple[bool, Any]: - """Inspects the given value and returns (is_overwrite, overwrite_value).""" - 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 BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): +class BinaryOperatorAggregate(AggregateChannel[Value], Generic[Value]): """Stores the result of applying a binary operator to the current value and each new value. ```python @@ -46,26 +23,16 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): total = Channels.BinaryOperatorAggregate(int, operator.add) ``` + + Equivalent to `AggregateChannel(operator, typ=typ, snapshot_frequency=1)`. + Preserved as a distinct subclass so existing `isinstance(x, BinaryOperatorAggregate)` + checks and `_is_field_binop` detection continue to work. New code should + prefer `AggregateChannel` directly — especially when a non-unit + `snapshot_frequency` is wanted. """ - __slots__ = ("value", "operator") - def __init__(self, typ: type[Value], operator: Callable[[Value, Value], Value]): - super().__init__(typ) - self.operator = operator - # special forms from typing or collections.abc are not instantiable - # so we need to replace them with their concrete counterparts - typ = _strip_extras(typ) - if typ in (collections.abc.Sequence, collections.abc.MutableSequence): - typ = list - if typ in (collections.abc.Set, collections.abc.MutableSet): - typ = set - if typ in (collections.abc.Mapping, collections.abc.MutableMapping): - typ = dict - try: - self.value = typ() - except Exception: - self.value = MISSING + super().__init__(operator, typ=typ, snapshot_frequency=1) def __eq__(self, value: object) -> bool: return isinstance(value, BinaryOperatorAggregate) and ( @@ -74,61 +41,3 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): and self.operator.__name__ != "" else True ) - - @property - def ValueType(self) -> type[Value]: - """The type of the value stored in the channel.""" - return self.typ - - @property - def UpdateType(self) -> type[Value]: - """The type of the update received by the channel.""" - return self.typ - - def copy(self) -> Self: - """Return a copy of the channel.""" - empty = self.__class__(self.typ, self.operator) - empty.key = self.key - empty.value = self.value - return empty - - def from_checkpoint(self, checkpoint: Value) -> Self: - empty = self.__class__(self.typ, self.operator) - empty.key = self.key - if checkpoint is not MISSING: - empty.value = checkpoint - return empty - - def update(self, values: Sequence[Value]) -> bool: - if not values: - return False - if self.value is MISSING: - self.value = values[0] - values = values[1:] - seen_overwrite: bool = False - for value in values: - is_overwrite, overwrite_value = _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) - self.value = overwrite_value - seen_overwrite = True - continue - if not seen_overwrite: - self.value = self.operator(self.value, value) - return True - - def get(self) -> Value: - if self.value is MISSING: - raise EmptyChannelError() - return self.value - - def is_available(self) -> bool: - return self.value is not MISSING - - def checkpoint(self) -> Value: - return self.value diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 9573e63ae..4149f717b 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -47,7 +47,7 @@ from langgraph._internal._fields import ( from langgraph._internal._pydantic import create_model from langgraph._internal._runnable import coerce_to_runnable from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs -from langgraph.channels._delta import DeltaChannel +from langgraph.channels.aggregate import AggregateChannel from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate, _strip_extras from langgraph.channels.ephemeral_value import EphemeralValue @@ -1670,7 +1670,15 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None: # Search through all annotated medata to find channel annotations for item in meta: if isinstance(item, BaseChannel): - if isinstance(item, DeltaChannel) and hasattr(typ, "__origin__"): + # AggregateChannel instances without an explicit `typ` arg + # inherit the value type from the outer `Annotated[...]`. + # BinaryOperatorAggregate (a subclass) always sets typ + # explicitly, so _typ_provided is True and we skip inference. + if ( + isinstance(item, AggregateChannel) + and not item._typ_provided + and hasattr(typ, "__origin__") + ): origin = typ.__origin__ # Unwrap parameterized Required[X]/NotRequired[X] to X # (e.g. Annotated[NotRequired[dict[...]], ...]). diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 51350071d..680524fbd 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -8,7 +8,7 @@ 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.aggregate import AggregateChannel from langgraph.channels.base import BaseChannel from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec @@ -34,7 +34,13 @@ 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. + + For `AggregateChannel` spec with `snapshot_frequency != 1`, the stored + blob alternates between the full value and `DELTA_SENTINEL` based on + `is_snapshot_step(step)`. Non-snapshot steps store the sentinel; the + value is reconstructed from ancestor writes at read time. + """ ts = datetime.now(timezone.utc).isoformat() if channels is None: values = checkpoint["channel_values"] @@ -43,7 +49,11 @@ def create_checkpoint( for k in channels: if k not in checkpoint["channel_versions"]: continue - v = channels[k].checkpoint() + ch = channels[k] + if isinstance(ch, AggregateChannel) and not ch.is_snapshot_step(step): + values[k] = DELTA_SENTINEL + continue + v = ch.checkpoint() if v is not MISSING: values[k] = v return Checkpoint( @@ -57,6 +67,16 @@ def create_checkpoint( ) +def _needs_replay(spec: BaseChannel, stored: object) -> bool: + """True if `spec` is a delta-mode AggregateChannel and the stored + blob is empty/sentinel, requiring an ancestor walk to reconstruct.""" + if not isinstance(spec, AggregateChannel): + return False + if spec.snapshot_frequency == 1: + return False + return stored is MISSING or stored is DELTA_SENTINEL + + def channels_from_checkpoint( specs: Mapping[str, BaseChannel | ManagedValueSpec], checkpoint: Checkpoint, @@ -69,12 +89,13 @@ 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. + `AggregateChannel` with `snapshot_frequency != 1` 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. Without them (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,15 @@ 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). 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 +144,11 @@ 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: 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_channels.py b/libs/langgraph/tests/test_channels.py index 1aed2333f..a81fc1ee4 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -6,7 +6,6 @@ from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.base import DELTA_SENTINEL from langgraph._internal._typing import MISSING -from langgraph.channels._delta import DeltaChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -14,6 +13,12 @@ from langgraph.channels.untracked_value import UntrackedValue from langgraph.errors import EmptyChannelError, InvalidUpdateError from langgraph.graph.message import add_messages +import math as _math_compat +from langgraph.channels.aggregate import AggregateChannel as _AggregateChannel_compat +def DeltaChannel(op): + return _AggregateChannel_compat(op, snapshot_frequency=_math_compat.inf) + + pytestmark = pytest.mark.anyio @@ -127,7 +132,6 @@ def test_delta_channel_basic_two_steps() -> None: from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.channels._delta import DeltaChannel from langgraph.graph.message import add_messages ch = DeltaChannel(add_messages).from_checkpoint(MISSING) @@ -152,7 +156,6 @@ def test_delta_channel_from_checkpoint_writes_list() -> None: """replay_writes on a fresh channel replays through the operator.""" from langchain_core.messages import AIMessage, HumanMessage - from langgraph.channels._delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) @@ -174,7 +177,6 @@ def test_delta_channel_from_checkpoint_writes_list() -> None: def test_delta_channel_from_checkpoint_backwards_compat() -> None: from langchain_core.messages import HumanMessage - from langgraph.channels._delta import DeltaChannel from langgraph.graph.message import add_messages # Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat @@ -188,7 +190,6 @@ def test_delta_channel_overwrite() -> None: from langchain_core.messages import HumanMessage from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.channels._delta import DeltaChannel from langgraph.graph.message import add_messages from langgraph.types import Overwrite @@ -207,7 +208,6 @@ 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.channels._delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) @@ -241,7 +241,6 @@ 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.channels._delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) @@ -270,7 +269,6 @@ def test_delta_channel_checkpoint_returns_sentinel() -> None: """checkpoint() always returns DELTA_SENTINEL regardless of state.""" from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.channels._delta import DeltaChannel from langgraph.graph.message import add_messages ch = DeltaChannel(add_messages).from_checkpoint(MISSING) @@ -290,7 +288,6 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None: 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 @@ -329,7 +326,6 @@ def _delta_channel_with_type(operator, typ): """Build a DeltaChannel with an explicit type via the Annotated injection path.""" from typing import Annotated - from langgraph.channels._delta import DeltaChannel from langgraph.graph.state import _get_channel return _get_channel("_test", Annotated[typ, DeltaChannel(operator)]) @@ -464,7 +460,6 @@ def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None: from typing_extensions import NotRequired - from langgraph.channels._delta import DeltaChannel from langgraph.graph.state import _get_channel def merge_dicts(left: dict | None, right: dict) -> dict: @@ -494,7 +489,6 @@ def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None: from langgraph.checkpoint.memory import InMemorySaver from typing_extensions import TypedDict - from langgraph.channels._delta import DeltaChannel from langgraph.graph import START, StateGraph def merge_files(left: dict | None, right: dict) -> dict: diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index bc393d804..07d0bcead 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -26,10 +26,15 @@ 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.graph import END, StateGraph from langgraph.graph.message import add_messages +import math as _math_compat +from langgraph.channels.aggregate import AggregateChannel as _AggregateChannel_compat +def DeltaChannel(op): + return _AggregateChannel_compat(op, snapshot_frequency=_math_compat.inf) + + try: from langgraph.checkpoint.sqlite import SqliteSaver diff --git a/libs/langgraph/tests/test_delta_channel_migration.py b/libs/langgraph/tests/test_delta_channel_migration.py index fb107395a..e6703f5d1 100644 --- a/libs/langgraph/tests/test_delta_channel_migration.py +++ b/libs/langgraph/tests/test_delta_channel_migration.py @@ -49,10 +49,15 @@ import pytest from langgraph.checkpoint.memory import InMemorySaver from typing_extensions import TypedDict -from langgraph.channels._delta import DeltaChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.graph import END, START, StateGraph +import math as _math_compat +from langgraph.channels.aggregate import AggregateChannel as _AggregateChannel_compat +def DeltaChannel(op): + return _AggregateChannel_compat(op, snapshot_frequency=_math_compat.inf) + + pytestmark = pytest.mark.anyio diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 474e08c51..36db4fc7f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -79,6 +79,14 @@ from tests.messages import ( _AnyIdToolMessage, ) +import math as _math_compat +from langgraph.channels.aggregate import AggregateChannel as _AggregateChannel_compat + + +def DeltaChannel(op): + return _AggregateChannel_compat(op, snapshot_frequency=_math_compat.inf) + + pytestmark = pytest.mark.anyio logger = logging.getLogger(__name__) @@ -9407,7 +9415,6 @@ async def test_delta_channel_end_to_end_inmemory() -> None: from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.memory import InMemorySaver - from langgraph.channels._delta import DeltaChannel from langgraph.graph import START, StateGraph from langgraph.graph.message import add_messages @@ -9449,7 +9456,6 @@ async def test_delta_channel_time_travel() -> None: from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.memory import InMemorySaver - from langgraph.channels._delta import DeltaChannel from langgraph.graph import START, StateGraph from langgraph.graph.message import add_messages @@ -9507,7 +9513,6 @@ async def test_delta_channel_remove_message_end_to_end() -> None: from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage from langgraph.checkpoint.memory import InMemorySaver - from langgraph.channels._delta import DeltaChannel from langgraph.graph import START, StateGraph from langgraph.graph.message import add_messages @@ -9554,7 +9559,6 @@ async def test_delta_channel_update_by_id_end_to_end() -> None: from langchain_core.messages import HumanMessage from langgraph.checkpoint.memory import InMemorySaver - from langgraph.channels._delta import DeltaChannel from langgraph.graph import START, StateGraph from langgraph.graph.message import add_messages