mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ee5b3fb4ca | ||
|
|
f247a1a647 |
@@ -0,0 +1,219 @@
|
||||
# AggregateChannel — unified fold-reducer channel with configurable snapshot cadence
|
||||
|
||||
**Status:** MVP scope approved. Implementation starting 2026-04-24.
|
||||
**Supersedes:** `langgraph.channels._delta.DeltaChannel` (experimental, private).
|
||||
**Branch:** `sr/even-better-writes-idea` (forked from `delta-channel-writes-based`).
|
||||
|
||||
## Problem
|
||||
|
||||
The experimental `DeltaChannel` (introduced earlier on this branch) stores a
|
||||
sentinel in every checkpoint and reconstructs state by walking ancestor
|
||||
writes. It eliminates O(N²) blob growth for append-style reducers on long
|
||||
threads, but read cost now scales O(N) with thread depth — every load
|
||||
replays every write since the start of the thread.
|
||||
|
||||
`BinaryOperatorAggregate` is the opposite extreme: always snapshots the full
|
||||
value every step. Zero replay cost at read time, but O(N²) storage on
|
||||
append-heavy workloads.
|
||||
|
||||
These are endpoints of the same axis. A single channel class parameterised
|
||||
on snapshot cadence covers both — plus every intermediate point.
|
||||
|
||||
The concrete pain that surfaced this design: deep-agent workloads at
|
||||
200+ turns pay O(200) replay per read under `DeltaChannel`. A
|
||||
`snapshot_frequency` knob bounds that to O(snapshot_frequency) regardless
|
||||
of thread depth.
|
||||
|
||||
## MVP scope (this PR)
|
||||
|
||||
1. **New class `AggregateChannel`** at `libs/langgraph/langgraph/channels/aggregate.py`:
|
||||
|
||||
```python
|
||||
class AggregateChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Value, Value], Value],
|
||||
*,
|
||||
snapshot_frequency: int | float = 1,
|
||||
typ: type[Value] | None = None,
|
||||
): ...
|
||||
```
|
||||
|
||||
- `snapshot_frequency=1` (default): full snapshot every step. Equivalent
|
||||
to today's `BinaryOperatorAggregate`.
|
||||
- `snapshot_frequency=N` (integer > 1): full snapshot every Nth step;
|
||||
sentinel on other steps.
|
||||
- `snapshot_frequency=math.inf`: never snapshot. Equivalent to today's
|
||||
`DeltaChannel`.
|
||||
- `typ` inferred from `Annotated[...]` via `_strip_extras` when used in
|
||||
a `TypedDict` state schema; kwarg is the escape hatch for imperative
|
||||
constructions.
|
||||
|
||||
2. **Rewire `BinaryOperatorAggregate` as a subclass** of `AggregateChannel`
|
||||
with `snapshot_frequency=1` hard-coded. Preserves
|
||||
`isinstance(x, BinaryOperatorAggregate)` for any existing user code and
|
||||
`_is_field_binop` detection in `graph/state.py`.
|
||||
|
||||
```python
|
||||
class BinaryOperatorAggregate(AggregateChannel):
|
||||
def __init__(self, typ, operator):
|
||||
super().__init__(operator, typ=typ, snapshot_frequency=1)
|
||||
```
|
||||
|
||||
3. **Delete `langgraph/channels/_delta.py`** (`DeltaChannel`). It was
|
||||
experimental, private, underscored, and not re-exported — clean
|
||||
removal. Users migrate to `AggregateChannel(op, snapshot_frequency=math.inf)`.
|
||||
|
||||
4. **Step-aware `create_checkpoint`** at `libs/langgraph/langgraph/pregel/_checkpoint.py`:
|
||||
|
||||
`AggregateChannel` exposes a helper method:
|
||||
|
||||
```python
|
||||
def is_snapshot_step(self, step: int) -> bool:
|
||||
if self.snapshot_frequency == 1:
|
||||
return True
|
||||
if self.snapshot_frequency == math.inf:
|
||||
return False
|
||||
return step % self.snapshot_frequency == 0
|
||||
```
|
||||
|
||||
`create_checkpoint` calls this per channel. When it returns `False`,
|
||||
the row stores `DELTA_SENTINEL` for that channel; when `True`, it
|
||||
stores `ch.checkpoint()` as today. `step` is already an argument to
|
||||
`create_checkpoint` — no new threading. The explicit branches on
|
||||
`1` and `math.inf` avoid relying on `step % math.inf` NaN arithmetic
|
||||
and make the common cases (always snapshot / never snapshot) free of
|
||||
a modulo.
|
||||
|
||||
5. **Generalise `channels_from_checkpoint`**. The existing DeltaChannel-specific
|
||||
branch keys off `isinstance(spec, DeltaChannel)`; change to
|
||||
`isinstance(spec, AggregateChannel) and spec.snapshot_frequency != 1`.
|
||||
The existing "pre-delta seed terminator" walk already treats any
|
||||
non-sentinel ancestor blob as the base value and stops — no change
|
||||
needed to saver-side logic. A `snapshot_frequency=10` blob at step 50
|
||||
serves as a natural terminator for a walk starting at step 57.
|
||||
|
||||
6. **Saver API stays as-is.** `_get_channel_writes_history` (private,
|
||||
underscored) continues to be the reconstruction hook. The broader
|
||||
refactor (`walk_writes` + `put_channel_snapshot`, batched multi-channel
|
||||
walks) is deferred to a follow-up PR that can benchmark its own win
|
||||
independently.
|
||||
|
||||
## Explicitly deferred (documented, not implemented here)
|
||||
|
||||
Each of the below lands as its own PR on top of this MVP.
|
||||
|
||||
- **`coalesce=` kwarg on `AggregateChannel`.** Batch-shape reducer for users
|
||||
who need to see all of a step's writes at once (non-binary-foldable
|
||||
reducers: median, priority-pick, dedup-across-writes). Additive to the
|
||||
existing `operator` kwarg; exactly one of the two must be provided.
|
||||
- **Saver API boundary refactor.** Rename `_get_channel_writes_history` →
|
||||
`walk_writes(config, *, channels=None)`. Move the `DELTA_SENTINEL`
|
||||
terminator check from saver-side to pregel-side. Saver becomes a pure
|
||||
storage primitive (aligns with every event-sourced system surveyed:
|
||||
Akka Persistence, EventStoreDB, Kafka Streams, Postgres logical
|
||||
replication, Firestore). Research memo captured in brainstorming session.
|
||||
- **Batched multi-channel walks.** One ancestor-walk query per read
|
||||
regardless of how many `AggregateChannel` channels need hydration.
|
||||
Today each channel triggers its own walk; deep-agent graphs with 7 state
|
||||
channels pay 7× the latency.
|
||||
- **`put_channel_snapshot` saver hook.** Opportunistic/manual compaction of
|
||||
a pure-delta (`snapshot_frequency=math.inf`) thread by retroactively
|
||||
promoting a sentinel row to a full blob. Separate from the write hot path.
|
||||
- **ShallowPostgresSaver compat.** `snapshot_frequency > 1` is fundamentally
|
||||
incompatible with shallow savers (no parent chain → nowhere to walk).
|
||||
Detect at attach time and error loudly. The existing `DeltaChannel` has
|
||||
the same silent incompatibility today; make it explicit in the same
|
||||
pass.
|
||||
- **Option A — channel_versions / versions_seen delta-encoding.** Documented
|
||||
in `notes/delta_checkpoint_rows.md`. 60% win on the checkpoint row table,
|
||||
reuses the same parent-walk machinery. Requires `Checkpoint.v` bump
|
||||
(4 → 5), so wants to land on top of the saver API refactor, not stacked
|
||||
with this MVP.
|
||||
|
||||
## Key design decisions and why
|
||||
|
||||
- **`operator`-only MVP, `coalesce` deferred.** The deepagents workload
|
||||
uses `add_messages`-shape reducers (binary-foldable). Shipping
|
||||
`coalesce=` now expands the API surface before we've validated that
|
||||
the snapshot-cadence half works on a real workload. `coalesce=` is
|
||||
additive and can land without breaking anyone.
|
||||
- **Subclass, not alias.** `BinaryOperatorAggregate` is imported and
|
||||
instantiated directly in at least `libs/langgraph/langgraph/graph/state.py:1711`
|
||||
(`_is_field_binop`). A factory function breaks `isinstance`; a subclass
|
||||
doesn't.
|
||||
- **`snapshot_frequency`, not `snapshot_every`.** Chosen by user preference;
|
||||
semantically identical (integer period, default 1).
|
||||
- **No saver API change in MVP.** The saver's existing
|
||||
`_get_channel_writes_history` contract is sufficient for the cadence
|
||||
knob to work. Deferring the saver refactor lets this PR land
|
||||
independently and the refactor benchmark against a stable baseline.
|
||||
- **No benchmark harness in this PR.** Benchmarking happens externally
|
||||
against deepagents.
|
||||
- **DELTA_SENTINEL keeps its name.** Even though the class renames to
|
||||
`AggregateChannel`, the sentinel itself is still "this row represents a
|
||||
delta from ancestors" — the name is accurate. Rename could happen in a
|
||||
later cleanup if desired but isn't in scope.
|
||||
|
||||
## Migration semantics
|
||||
|
||||
- **Existing threads with `BinaryOperatorAggregate`** continue to work
|
||||
unchanged — they're now instances of `AggregateChannel` with
|
||||
`snapshot_frequency=1`, and the runtime code paths are identical.
|
||||
- **Switching `snapshot_frequency` mid-thread** (e.g. user bumps
|
||||
`snapshot_frequency=1` → `10` on an existing thread): pre-change
|
||||
checkpoints have full blobs; they act as natural walk terminators for
|
||||
post-change reads. No explicit migration step. No data loss.
|
||||
- **Reverse migration** (`snapshot_frequency=10` → `1`): next write produces
|
||||
a full blob. Reads at ancestors still find the right terminator. Safe.
|
||||
- **Switching `snapshot_frequency=math.inf` → any finite value:**
|
||||
next snapshot-step writes a full blob that closes all prior sentinel-only
|
||||
ancestry. Walks from that point forward stop at the new base rather
|
||||
than walking to the root.
|
||||
- **External users who imported `langgraph.channels._delta.DeltaChannel`**:
|
||||
`ImportError` at upgrade time. Underscored + experimental + docstring
|
||||
says "subject to change or removal without notice" — documented
|
||||
breakage; migration is a one-line swap.
|
||||
|
||||
## File-level change list
|
||||
|
||||
**New:**
|
||||
- `libs/langgraph/langgraph/channels/aggregate.py` — `AggregateChannel` class
|
||||
|
||||
**Modified:**
|
||||
- `libs/langgraph/langgraph/channels/binop.py` — `BinaryOperatorAggregate`
|
||||
becomes subclass of `AggregateChannel`; reducer logic moves to base.
|
||||
- `libs/langgraph/langgraph/channels/__init__.py` — export
|
||||
`AggregateChannel`.
|
||||
- `libs/langgraph/langgraph/pregel/_checkpoint.py`:
|
||||
- `create_checkpoint`: step-aware sentinel vs blob decision.
|
||||
- `channels_from_checkpoint` / `achannels_from_checkpoint`: key off
|
||||
`AggregateChannel` instead of `DeltaChannel`.
|
||||
- `DeltaChannel` import removed.
|
||||
- `libs/langgraph/langgraph/graph/state.py` — `_is_field_binop` continues
|
||||
to work unchanged (subclass relationship preserves detection).
|
||||
|
||||
**Deleted:**
|
||||
- `libs/langgraph/langgraph/channels/_delta.py`
|
||||
|
||||
**Tests updated:**
|
||||
- Any test importing `DeltaChannel` from `langgraph.channels._delta` —
|
||||
switch to `AggregateChannel(op, snapshot_frequency=math.inf)`.
|
||||
- Add: parity test for `snapshot_frequency=1` vs today's `BinaryOperatorAggregate`
|
||||
on the same workload.
|
||||
- Add: cadence test at `snapshot_frequency=10` on a 50-step thread —
|
||||
verify blobs land on steps 0/10/20/30/40/50, sentinels elsewhere, and
|
||||
reads at every step produce the same value as the all-snapshot baseline.
|
||||
- Add: mid-thread `snapshot_frequency` change — verify no data loss.
|
||||
|
||||
## Out of scope
|
||||
|
||||
- Performance benchmarking (user runs externally on deepagents).
|
||||
- Documentation/tutorial updates for the new knob (until MVP validates).
|
||||
- Any saver-side changes.
|
||||
- Any `Checkpoint.v` bump.
|
||||
- Public API promotion — `AggregateChannel` replaces a private
|
||||
experimental class; it's immediately public by virtue of living in
|
||||
`langgraph.channels`, but the `snapshot_frequency > 1` path inherits
|
||||
DeltaChannel's "experimental, validate on real workloads first"
|
||||
caveat until benchmark confirms it.
|
||||
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
MetadataInput = dict[str, Any] | None
|
||||
@@ -281,34 +281,22 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
collected: list[PendingWrite] = [] # newest first; reversed at the end
|
||||
for cid in ancestors:
|
||||
# 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)
|
||||
if seed_blob is not None and seed_blob[0] != "empty":
|
||||
blob_value = self.serde.loads_typed(seed_blob)
|
||||
if blob_value is not DELTA_SENTINEL:
|
||||
if isinstance(blob_value, _DeltaSnapshot):
|
||||
# Step-based snapshot: collect this ancestor's
|
||||
# pending_writes first (they encode the NEXT step's
|
||||
# transition, not subsumed by the snapshot blob).
|
||||
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()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
)
|
||||
# Pre-delta blob: subsumes this ancestor's writes.
|
||||
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)
|
||||
|
||||
@@ -377,13 +377,13 @@ async def test_get_checkpoint_no_channel_values(
|
||||
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
|
||||
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
|
||||
pytest.importorskip(
|
||||
"langgraph.channels.delta", reason="langgraph core not installed"
|
||||
"langgraph.channels._delta", reason="langgraph core not installed"
|
||||
)
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
@@ -28,7 +28,6 @@ from langgraph.checkpoint.serde.types import (
|
||||
RESUME,
|
||||
SCHEDULED,
|
||||
ChannelProtocol,
|
||||
_DeltaSnapshot,
|
||||
)
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
@@ -541,23 +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:
|
||||
if isinstance(ancestor_value, _DeltaSnapshot):
|
||||
# Step-based snapshot: the blob is state AT this ancestor,
|
||||
# but pending_writes encode the NEXT step's transition and
|
||||
# are NOT subsumed — collect them before terminating.
|
||||
if tup.pending_writes:
|
||||
for write in reversed(tup.pending_writes):
|
||||
if write[1] != channel:
|
||||
continue
|
||||
collected.append(write)
|
||||
# Pre-delta blob: subsumes its own writes — stop immediately.
|
||||
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.
|
||||
@@ -565,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)
|
||||
@@ -589,22 +586,17 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
tup = await self.aget_tuple(cursor_config)
|
||||
if tup is None:
|
||||
break
|
||||
# See sync variant for rationale.
|
||||
ancestor_value = tup.checkpoint["channel_values"].get(channel)
|
||||
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
|
||||
if isinstance(ancestor_value, _DeltaSnapshot):
|
||||
if tup.pending_writes:
|
||||
for write in reversed(tup.pending_writes):
|
||||
if write[1] != channel:
|
||||
continue
|
||||
collected.append(write)
|
||||
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)
|
||||
|
||||
@@ -27,7 +27,6 @@ from langgraph.checkpoint.base import (
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -162,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])
|
||||
@@ -186,46 +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:
|
||||
if isinstance(blob_value, _DeltaSnapshot):
|
||||
# Step-based snapshot: the blob is state AT this
|
||||
# ancestor, but the ancestor's pending_writes
|
||||
# encode the NEXT step's transition and are NOT
|
||||
# subsumed by the snapshot — collect them first.
|
||||
step_writes = self.writes.get(
|
||||
(thread_id, checkpoint_ns, cp_id), {}
|
||||
)
|
||||
for (_task_id, _idx), (
|
||||
tid,
|
||||
ch,
|
||||
serialized,
|
||||
_,
|
||||
) in sorted(step_writes.items(), reverse=True):
|
||||
if ch != channel:
|
||||
continue
|
||||
collected.append(
|
||||
(tid, ch, self.serde.loads_typed(serialized))
|
||||
)
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
)
|
||||
# Pre-delta blob: state AT this ancestor already
|
||||
# subsumes its pending_writes — skip 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)
|
||||
|
||||
@@ -525,101 +498,6 @@ class InMemorySaver(
|
||||
task_path,
|
||||
)
|
||||
|
||||
def prune(
|
||||
self,
|
||||
thread_ids: Sequence[str],
|
||||
*,
|
||||
strategy: str = "keep_latest",
|
||||
) -> None:
|
||||
"""Prune checkpoints for the given threads.
|
||||
|
||||
For DeltaChannel channels, a checkpoint is only deleted if the walk
|
||||
from the latest checkpoint would not need to traverse it — i.e., a
|
||||
`_DeltaSnapshot` blob exists in the kept ancestry that covers all
|
||||
sentinel channels. Checkpoints that are still in the active walk
|
||||
chain (because no snapshot has been taken yet, e.g. with
|
||||
`snapshot_frequency=None`) are retained.
|
||||
|
||||
Args:
|
||||
thread_ids: Thread IDs to prune.
|
||||
strategy: ``"keep_latest"`` keeps only the most recent checkpoint
|
||||
per namespace; ``"delete"`` removes all checkpoints.
|
||||
"""
|
||||
for thread_id in thread_ids:
|
||||
if strategy == "delete":
|
||||
self.delete_thread(thread_id)
|
||||
continue
|
||||
|
||||
if strategy != "keep_latest":
|
||||
raise ValueError(
|
||||
f"Unknown pruning strategy {strategy!r}. "
|
||||
"Expected 'keep_latest' or 'delete'."
|
||||
)
|
||||
|
||||
for checkpoint_ns, ns_storage in list(
|
||||
self.storage.get(thread_id, {}).items()
|
||||
):
|
||||
if not ns_storage:
|
||||
continue
|
||||
|
||||
# Latest checkpoint (uuid6 IDs are lexicographically monotonic)
|
||||
latest_id = max(ns_storage.keys())
|
||||
latest_data, _, _ = ns_storage[latest_id]
|
||||
latest_cp = self.serde.loads_typed(latest_data)
|
||||
|
||||
# Which channels in the latest checkpoint still have sentinels?
|
||||
sentinel_channels: set[str] = set()
|
||||
for ch, ver in latest_cp.get("channel_versions", {}).items():
|
||||
blob = self.blobs.get((thread_id, checkpoint_ns, ch, ver))
|
||||
if blob is not None and blob[0] != "empty":
|
||||
if self.serde.loads_typed(blob) is DELTA_SENTINEL:
|
||||
sentinel_channels.add(ch)
|
||||
|
||||
# Walk the parent chain to find the oldest ancestor still needed.
|
||||
# We stop (and mark "safe to prune before here") when all
|
||||
# sentinel channels are covered by a non-sentinel blob.
|
||||
required_ids: set[str] = {latest_id}
|
||||
if sentinel_channels:
|
||||
_, _, parent_id = ns_storage[latest_id]
|
||||
remaining = set(sentinel_channels)
|
||||
while parent_id is not None and remaining:
|
||||
entry = ns_storage.get(parent_id)
|
||||
if entry is None:
|
||||
break
|
||||
required_ids.add(parent_id)
|
||||
cp_data, _, grandparent_id = entry
|
||||
cp = self.serde.loads_typed(cp_data)
|
||||
resolved: set[str] = set()
|
||||
for ch in remaining:
|
||||
ver = cp.get("channel_versions", {}).get(ch)
|
||||
if ver is None:
|
||||
continue
|
||||
blob = self.blobs.get((thread_id, checkpoint_ns, ch, ver))
|
||||
if blob is not None and blob[0] != "empty":
|
||||
if self.serde.loads_typed(blob) is not DELTA_SENTINEL:
|
||||
resolved.add(ch)
|
||||
remaining -= resolved
|
||||
parent_id = grandparent_id
|
||||
|
||||
# Delete everything outside the required set
|
||||
for cp_id in list(ns_storage.keys()):
|
||||
if cp_id in required_ids:
|
||||
continue
|
||||
cp_data, _, _ = ns_storage.pop(cp_id)
|
||||
self.writes.pop((thread_id, checkpoint_ns, cp_id), None)
|
||||
|
||||
# Clean up blobs no longer referenced by any kept checkpoint
|
||||
live: set[tuple[str, str, str, Any]] = set()
|
||||
for cp_data, _, _ in ns_storage.values():
|
||||
cp = self.serde.loads_typed(cp_data)
|
||||
for ch, ver in cp.get("channel_versions", {}).items():
|
||||
live.add((thread_id, checkpoint_ns, ch, ver))
|
||||
for key in [
|
||||
k for k in self.blobs if k[:2] == (thread_id, checkpoint_ns)
|
||||
]:
|
||||
if key not in live:
|
||||
del self.blobs[key]
|
||||
|
||||
def delete_thread(self, thread_id: str) -> None:
|
||||
"""Delete all checkpoints and writes associated with a thread ID.
|
||||
|
||||
|
||||
@@ -33,12 +33,7 @@ from langchain_core.load.load import Reviver
|
||||
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
DELTA_SENTINEL,
|
||||
SendProtocol,
|
||||
_DeltaSentinel,
|
||||
_DeltaSnapshot,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import DELTA_SENTINEL, SendProtocol
|
||||
from langgraph.store.base import Item
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -256,6 +251,8 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
|
||||
if obj is None:
|
||||
return "null", EMPTY_BYTES
|
||||
elif obj is DELTA_SENTINEL:
|
||||
return "delta", EMPTY_BYTES
|
||||
elif isinstance(obj, bytes):
|
||||
return "bytes", obj
|
||||
elif isinstance(obj, bytearray):
|
||||
@@ -282,6 +279,8 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
elif type_ == "delta":
|
||||
return DELTA_SENTINEL
|
||||
elif self.pickle_fallback and type_ == "pickle":
|
||||
return pickle.loads(data_)
|
||||
else:
|
||||
@@ -297,16 +296,10 @@ EXT_METHOD_SINGLE_ARG = 3
|
||||
EXT_PYDANTIC_V1 = 4
|
||||
EXT_PYDANTIC_V2 = 5
|
||||
EXT_NUMPY_ARRAY = 6
|
||||
EXT_DELTA_SNAPSHOT = 7
|
||||
EXT_DELTA_SENTINEL = 8
|
||||
|
||||
|
||||
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
if isinstance(obj, _DeltaSentinel):
|
||||
return ormsgpack.Ext(EXT_DELTA_SENTINEL, b"")
|
||||
elif isinstance(obj, _DeltaSnapshot):
|
||||
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
|
||||
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
||||
if hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
||||
return ormsgpack.Ext(
|
||||
EXT_PYDANTIC_V2,
|
||||
_msgpack_enc(
|
||||
@@ -620,15 +613,7 @@ def _create_msgpack_ext_hook(
|
||||
return False
|
||||
|
||||
def ext_hook(code: int, data: bytes) -> Any:
|
||||
if code == EXT_DELTA_SENTINEL:
|
||||
return DELTA_SENTINEL
|
||||
elif code == EXT_DELTA_SNAPSHOT:
|
||||
return _DeltaSnapshot(
|
||||
ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
)
|
||||
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import (
|
||||
Any,
|
||||
NamedTuple,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
runtime_checkable,
|
||||
@@ -34,20 +33,6 @@ class _DeltaSentinel:
|
||||
DELTA_SENTINEL = _DeltaSentinel()
|
||||
|
||||
|
||||
class _DeltaSnapshot(NamedTuple):
|
||||
"""Snapshot blob for a DeltaChannel with finite snapshot_frequency.
|
||||
|
||||
Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code.
|
||||
The ancestor walk in `_get_channel_writes_history` terminates when it
|
||||
encounters this type (any non-sentinel blob stops the walk).
|
||||
|
||||
`from_checkpoint` reconstructs the channel value directly from `.value`
|
||||
without replaying writes — the snapshot IS the accumulated state.
|
||||
"""
|
||||
|
||||
value: Any
|
||||
|
||||
|
||||
Value = TypeVar("Value", covariant=True)
|
||||
Update = TypeVar("Update", contravariant=True)
|
||||
C = TypeVar("C")
|
||||
|
||||
@@ -663,148 +663,3 @@ class TestPreDeltaBlobTerminator:
|
||||
assert "PRE-DELTA-WRITE" not in values
|
||||
# And the pending write at the target is never folded in.
|
||||
assert "PENDING-AT-TARGET" not in values
|
||||
|
||||
|
||||
class TestInMemorySaverPrune:
|
||||
"""Tests for InMemorySaver.prune with DeltaChannel awareness."""
|
||||
|
||||
def _build_chain(
|
||||
self,
|
||||
saver: InMemorySaver,
|
||||
thread_id: str,
|
||||
ns: str,
|
||||
channel: str,
|
||||
n: int,
|
||||
*,
|
||||
snapshot_at: set[int] | None = None,
|
||||
) -> list[str]:
|
||||
"""Build a chain of n checkpoints with DELTA_SENTINEL blobs.
|
||||
|
||||
If snapshot_at is provided, writes a _DeltaSnapshot blob at those steps.
|
||||
Returns list of checkpoint IDs in order (oldest first).
|
||||
"""
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
serde = saver.serde
|
||||
cp_ids = []
|
||||
parent_id = None
|
||||
ver_base = "0000000000000000000000000000000{i}.0000000000000000"
|
||||
|
||||
for i in range(n):
|
||||
cp_id = f"cp{i:04d}"
|
||||
ver = ver_base.format(i=i)
|
||||
cp = empty_checkpoint()
|
||||
cp["id"] = cp_id
|
||||
cp["channel_versions"][channel] = ver
|
||||
|
||||
if snapshot_at and i in snapshot_at:
|
||||
blob = serde.dumps_typed(_DeltaSnapshot(value=[f"msg{i}"]))
|
||||
else:
|
||||
blob = serde.dumps_typed(DELTA_SENTINEL)
|
||||
|
||||
saver.blobs[(thread_id, ns, channel, ver)] = blob
|
||||
saver.storage[thread_id][ns][cp_id] = (
|
||||
serde.dumps_typed(cp),
|
||||
serde.dumps_typed({}),
|
||||
parent_id,
|
||||
)
|
||||
# Add a dummy write for this checkpoint
|
||||
saver.writes[(thread_id, ns, cp_id)][("task", i)] = (
|
||||
"task",
|
||||
channel,
|
||||
serde.dumps_typed(f"write{i}"),
|
||||
"",
|
||||
)
|
||||
cp_ids.append(cp_id)
|
||||
parent_id = cp_id
|
||||
|
||||
return cp_ids
|
||||
|
||||
def test_prune_pure_delta_keeps_all(self) -> None:
|
||||
"""With no snapshots, all checkpoints are required for reconstruction."""
|
||||
saver = InMemorySaver()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
cp_ids = self._build_chain(saver, thread_id, ns, channel, n=5)
|
||||
|
||||
saver.prune([thread_id], strategy="keep_latest")
|
||||
|
||||
# All checkpoints must be retained (walk needs the full chain)
|
||||
remaining = set(saver.storage[thread_id][ns].keys())
|
||||
assert remaining == set(cp_ids)
|
||||
|
||||
def test_prune_with_snapshot_removes_pre_snapshot_checkpoints(self) -> None:
|
||||
"""Checkpoints older than the nearest snapshot can be safely pruned."""
|
||||
saver = InMemorySaver()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
# Snapshot at step 2; steps 3 and 4 are sentinels
|
||||
cp_ids = self._build_chain(saver, thread_id, ns, channel, n=5, snapshot_at={2})
|
||||
|
||||
saver.prune([thread_id], strategy="keep_latest")
|
||||
|
||||
remaining = set(saver.storage[thread_id][ns].keys())
|
||||
# cp0, cp1 (before snapshot) must be gone; cp2, cp3, cp4 must remain
|
||||
assert cp_ids[0] not in remaining # pre-snapshot
|
||||
assert cp_ids[1] not in remaining # pre-snapshot
|
||||
assert cp_ids[2] in remaining # the snapshot itself
|
||||
assert cp_ids[3] in remaining # sentinel after snapshot
|
||||
assert cp_ids[4] in remaining # latest
|
||||
|
||||
def test_prune_removes_orphaned_blobs(self) -> None:
|
||||
"""Blob entries for pruned checkpoints are cleaned up."""
|
||||
saver = InMemorySaver()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
self._build_chain(saver, thread_id, ns, channel, n=4, snapshot_at={1})
|
||||
saver.prune([thread_id], strategy="keep_latest")
|
||||
|
||||
# cp0 blob should be gone (pruned); cp1, cp2, cp3 blobs remain
|
||||
assert (
|
||||
thread_id,
|
||||
ns,
|
||||
channel,
|
||||
f"0000000000000000000000000000000{0}.0000000000000000",
|
||||
) not in saver.blobs
|
||||
for i in range(1, 4):
|
||||
ver = f"0000000000000000000000000000000{i}.0000000000000000"
|
||||
assert (thread_id, ns, channel, ver) in saver.blobs
|
||||
|
||||
def test_prune_removes_writes_for_pruned_checkpoints(self) -> None:
|
||||
"""checkpoint_writes for pruned checkpoints are deleted."""
|
||||
saver = InMemorySaver()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
cp_ids = self._build_chain(saver, thread_id, ns, channel, n=4, snapshot_at={1})
|
||||
|
||||
saver.prune([thread_id], strategy="keep_latest")
|
||||
|
||||
# writes for cp0 must be gone
|
||||
assert (thread_id, ns, cp_ids[0]) not in saver.writes
|
||||
# writes for cp1+ must remain (they're in the walk chain)
|
||||
for cp_id in cp_ids[1:]:
|
||||
assert (thread_id, ns, cp_id) in saver.writes
|
||||
|
||||
def test_prune_delete_strategy_removes_everything(self) -> None:
|
||||
"""strategy='delete' removes all checkpoints for the thread."""
|
||||
saver = InMemorySaver()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
self._build_chain(saver, thread_id, ns, channel, n=3)
|
||||
|
||||
saver.prune([thread_id], strategy="delete")
|
||||
|
||||
assert not saver.storage.get(thread_id, {}).get(ns)
|
||||
assert not any(k[0] == thread_id for k in saver.writes)
|
||||
assert not any(k[0] == thread_id for k in saver.blobs)
|
||||
|
||||
def test_prune_non_delta_channel_always_pruneable(self) -> None:
|
||||
"""A channel with full snapshot blobs (no sentinels) allows full prune."""
|
||||
|
||||
saver = InMemorySaver()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
# All snapshots, no sentinels
|
||||
cp_ids = self._build_chain(
|
||||
saver, thread_id, ns, channel, n=4, snapshot_at={0, 1, 2, 3}
|
||||
)
|
||||
|
||||
saver.prune([thread_id], strategy="keep_latest")
|
||||
|
||||
remaining = set(saver.storage[thread_id][ns].keys())
|
||||
# Only the latest checkpoint is needed (all blobs are snapshots)
|
||||
assert remaining == {cp_ids[-1]}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
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
|
||||
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,8 +20,8 @@ __all__ = (
|
||||
"LastValueAfterFinish",
|
||||
"UntrackedValue",
|
||||
"EphemeralValue",
|
||||
"AggregateChannel",
|
||||
"BinaryOperatorAggregate",
|
||||
"DeltaChannel",
|
||||
"NamedBarrierValue",
|
||||
"NamedBarrierValueAfterFinish",
|
||||
# topics
|
||||
|
||||
@@ -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__ != "<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:
|
||||
"""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
|
||||
@@ -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__ != "<lambda>"
|
||||
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
|
||||
|
||||
@@ -1,181 +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 langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
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=None` (default): pure delta — stores only
|
||||
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes.
|
||||
|
||||
`snapshot_frequency=N`: pregel's `create_checkpoint` writes a full
|
||||
`_DeltaSnapshot` blob every N steps (eagerly, even if the channel had
|
||||
no write that step). Reads walk at most N ancestor checkpoints before
|
||||
hitting the snapshot, bounding replay depth to N regardless of thread
|
||||
length.
|
||||
|
||||
Parameters:
|
||||
operator: Binary reducer `(Value, Value) -> Value`.
|
||||
snapshot_frequency: Every Nth pregel step writes a snapshot blob.
|
||||
`None` (default) = pure delta, never snapshot.
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "operator", "snapshot_frequency")
|
||||
value: Value | Any
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Any, Any], Any],
|
||||
*,
|
||||
snapshot_frequency: int | None = None,
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.snapshot_frequency = snapshot_frequency
|
||||
self.value: Any = []
|
||||
|
||||
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 is_snapshot_step(self, step: int) -> bool:
|
||||
"""True if pregel should write a snapshot blob at this step."""
|
||||
return (
|
||||
self.snapshot_frequency is not None and step % self.snapshot_frequency == 0
|
||||
)
|
||||
|
||||
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.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 _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 types (dispatched via serde ext code, not dict key inspection):
|
||||
* `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes.
|
||||
* `_DeltaSnapshot(value)`: restore value directly from snapshot.
|
||||
* plain value (migration from old BinOp blobs): use directly.
|
||||
"""
|
||||
new = self._clone_empty()
|
||||
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
|
||||
new.value = _empty(new.typ)
|
||||
elif isinstance(checkpoint, _DeltaSnapshot):
|
||||
new.value = checkpoint.value
|
||||
else:
|
||||
new.value = checkpoint
|
||||
return new
|
||||
|
||||
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
|
||||
"""Fold ancestor writes oldest→newest into current value."""
|
||||
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:
|
||||
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) -> 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: always `DELTA_SENTINEL`.
|
||||
|
||||
Snapshot decisions are made by `create_checkpoint` in pregel (which
|
||||
has the step number) via `is_snapshot_step`. `checkpoint()` is only
|
||||
called for non-snapshot steps or when no checkpointer is available.
|
||||
"""
|
||||
if self.value is MISSING:
|
||||
return MISSING
|
||||
return DELTA_SENTINEL
|
||||
@@ -47,9 +47,9 @@ 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.aggregate import AggregateChannel
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate, _strip_extras
|
||||
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 (
|
||||
@@ -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[...]], ...]).
|
||||
|
||||
@@ -210,7 +210,7 @@ def local_read(
|
||||
return values
|
||||
|
||||
|
||||
def increment(current: int | None, channel: None = None) -> int:
|
||||
def increment(current: int | None, channel: None) -> int:
|
||||
"""Default channel versioning function, increments the current int version."""
|
||||
return current + 1 if current is not None else 1
|
||||
|
||||
|
||||
@@ -1,23 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.aggregate import AggregateChannel
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
LATEST_VERSION = 4
|
||||
|
||||
GetNextVersion = Callable[[Any, None], Any]
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
@@ -37,63 +33,46 @@ def create_checkpoint(
|
||||
*,
|
||||
id: str | None = None,
|
||||
updated_channels: set[str] | None = None,
|
||||
get_next_version: GetNextVersion | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels.
|
||||
|
||||
For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a
|
||||
`_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor
|
||||
walk to at most N steps. Snapshots are eager: even if the channel had no
|
||||
write this step, a version bump is forced (via `get_next_version`) so the
|
||||
blob is stored by `put()`. Without `get_next_version` (e.g. static
|
||||
contexts), snapshot steps gracefully fall back to sentinel.
|
||||
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"]
|
||||
channel_versions = checkpoint["channel_versions"]
|
||||
else:
|
||||
values = {}
|
||||
channel_versions = dict(checkpoint["channel_versions"])
|
||||
for k in channels:
|
||||
if k not in channel_versions:
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
continue
|
||||
ch = channels[k]
|
||||
if (
|
||||
isinstance(ch, DeltaChannel)
|
||||
and ch.is_snapshot_step(step)
|
||||
and ch.is_available()
|
||||
):
|
||||
# Eager snapshot: bump version if not already written this step
|
||||
# so put() includes this channel in new_versions and stores blob.
|
||||
if get_next_version is not None and (
|
||||
updated_channels is None or k not in updated_channels
|
||||
):
|
||||
channel_versions[k] = get_next_version(channel_versions[k], None)
|
||||
values[k] = _DeltaSnapshot(ch.get())
|
||||
else:
|
||||
v = ch.checkpoint()
|
||||
if v is not MISSING:
|
||||
values[k] = v
|
||||
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(
|
||||
v=LATEST_VERSION,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=channel_versions,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
updated_channels=None if updated_channels is None else sorted(updated_channels),
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
`_DeltaSnapshot` blobs and plain values (migration) resolve directly via
|
||||
`from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay.
|
||||
"""
|
||||
if not isinstance(spec, DeltaChannel):
|
||||
"""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
|
||||
|
||||
@@ -108,10 +87,15 @@ def channels_from_checkpoint(
|
||||
"""Hydrate channels from a checkpoint.
|
||||
|
||||
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
|
||||
is sufficient. `DeltaChannel` is the exception: sentinel blobs require an
|
||||
ancestor walk via `saver._get_channel_writes_history`. The walk terminates
|
||||
at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain
|
||||
value, so read depth is bounded by `snapshot_frequency`.
|
||||
is sufficient — the stored value IS the reconstructed state.
|
||||
|
||||
`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] = {}
|
||||
@@ -126,7 +110,10 @@ def channels_from_checkpoint(
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
if _needs_replay(spec, stored) and saver is not None and config is not None:
|
||||
assert isinstance(spec, DeltaChannel)
|
||||
# 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)
|
||||
replay_ch = spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
@@ -158,7 +145,6 @@ async def achannels_from_checkpoint(
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
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)
|
||||
replay_ch = spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
|
||||
@@ -25,7 +25,6 @@ from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.cache.base import BaseCache
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
@@ -199,7 +198,6 @@ class PregelLoop:
|
||||
checkpoint_pending_writes: list[PendingWrite]
|
||||
checkpoint_previous_versions: dict[str, str | float | int]
|
||||
prev_checkpoint_config: RunnableConfig | None
|
||||
_pending_write_futs: list[concurrent.futures.Future]
|
||||
|
||||
status: Literal[
|
||||
"input",
|
||||
@@ -409,7 +407,7 @@ class PregelLoop:
|
||||
task = self.tasks.get(task_id)
|
||||
else:
|
||||
task = None
|
||||
fut = self.submit(
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
config,
|
||||
writes_to_save,
|
||||
@@ -417,13 +415,12 @@ class PregelLoop:
|
||||
task_path_str(task.path) if task else "",
|
||||
)
|
||||
else:
|
||||
fut = self.submit(
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
config,
|
||||
writes_to_save,
|
||||
task_id,
|
||||
)
|
||||
self._pending_write_futs.append(fut)
|
||||
# output writes
|
||||
if hasattr(self, "tasks"):
|
||||
self.output_writes(task_id, writes)
|
||||
@@ -894,9 +891,6 @@ class PregelLoop:
|
||||
self.step,
|
||||
id=self.checkpoint["id"] if exiting else None,
|
||||
updated_channels=self.updated_channels,
|
||||
get_next_version=self.checkpointer_get_next_version
|
||||
if do_checkpoint
|
||||
else None,
|
||||
)
|
||||
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
|
||||
if TASKS in self.checkpoint["channel_values"] and any(
|
||||
@@ -934,17 +928,6 @@ class PregelLoop:
|
||||
)
|
||||
self.checkpoint_previous_versions = channel_versions
|
||||
|
||||
# If the checkpoint has any DELTA_SENTINEL blobs, the sentinel is
|
||||
# only meaningful if checkpoint_writes are durable first. Flush
|
||||
# pending write futures synchronously before committing the blob so
|
||||
# we never end up with a sentinel blob backed by missing writes.
|
||||
if self._pending_write_futs and any(
|
||||
v is DELTA_SENTINEL for v in self.checkpoint["channel_values"].values()
|
||||
):
|
||||
for fut in self._pending_write_futs:
|
||||
fut.result()
|
||||
self._pending_write_futs.clear()
|
||||
|
||||
# save it, without blocking
|
||||
# if there's a previous checkpoint save in progress, wait for it
|
||||
# ensuring checkpointers receive checkpoints in order
|
||||
@@ -1289,7 +1272,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
if saved.pending_writes is not None
|
||||
else []
|
||||
)
|
||||
self._pending_write_futs = []
|
||||
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs,
|
||||
@@ -1495,7 +1477,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
if saved.pending_writes is not None
|
||||
else []
|
||||
)
|
||||
self._pending_write_futs = []
|
||||
self.submit = await self.stack.enter_async_context(
|
||||
AsyncBackgroundExecutor(self.config)
|
||||
)
|
||||
|
||||
@@ -7,13 +7,18 @@ from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.topic import Topic
|
||||
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
|
||||
|
||||
|
||||
@@ -275,124 +280,6 @@ def test_delta_channel_checkpoint_returns_sentinel() -> None:
|
||||
assert ch.checkpoint() is DELTA_SENTINEL
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_step_based() -> None:
|
||||
"""Snapshots fire on every Nth step regardless of whether the channel was written.
|
||||
|
||||
With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot
|
||||
blob — even if the channel had no write that step (eager snapshot). This
|
||||
bounds the ancestor walk to at most N steps on any read.
|
||||
"""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# snapshot_frequency=5: snapshot every 5 pregel steps
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_frequency=5)]
|
||||
other: str
|
||||
|
||||
def node_a(state: State) -> dict:
|
||||
# writes to messages
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
||||
|
||||
def node_b(state: State) -> dict:
|
||||
# writes ONLY to other, not messages — snapshot must still fire at step N
|
||||
return {"other": "y"}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("a", node_a)
|
||||
g.add_node("b", node_b)
|
||||
g.add_edge(START, "a")
|
||||
g.add_edge("a", "b")
|
||||
saver = InMemorySaver()
|
||||
graph = g.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
for i in range(6):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""},
|
||||
config,
|
||||
)
|
||||
|
||||
# Confirm at least one snapshot blob exists for messages
|
||||
msg_blob_values = [
|
||||
saver.serde.loads_typed((type_tag, blob))
|
||||
for k, (type_tag, blob) in saver.blobs.items()
|
||||
if k[2] == "messages" and type_tag == "msgpack" and blob
|
||||
]
|
||||
snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)]
|
||||
assert snapshots, "expected at least one _DeltaSnapshot blob for messages"
|
||||
|
||||
# Final state must be correct regardless of snapshot cadence
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 12 # 6 human + 6 AI
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_fires_even_when_not_written() -> None:
|
||||
"""Eager snapshot: _DeltaSnapshot stored at snapshot step even when the
|
||||
channel had no write that step (node_b doesn't touch messages).
|
||||
"""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_frequency=3)]
|
||||
tick: int
|
||||
|
||||
def writer(state: State) -> dict:
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
||||
|
||||
def ticker(state: State) -> dict:
|
||||
# never writes messages
|
||||
return {"tick": state["tick"] + 1}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("writer", writer)
|
||||
g.add_node("ticker", ticker)
|
||||
g.add_edge(START, "writer")
|
||||
g.add_edge("writer", "ticker")
|
||||
saver = InMemorySaver()
|
||||
graph = g.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
for i in range(5):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0},
|
||||
config,
|
||||
)
|
||||
|
||||
# Count distinct message channel blob versions
|
||||
msg_blobs = {
|
||||
k: saver.serde.loads_typed((t, b))
|
||||
for k, (t, b) in saver.blobs.items()
|
||||
if k[2] == "messages" and t == "msgpack" and b
|
||||
}
|
||||
snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)}
|
||||
# There must be snapshots (ticker steps are snapshot steps too)
|
||||
assert snapshots, (
|
||||
"eager snapshots must fire even on steps where messages wasn't written"
|
||||
)
|
||||
|
||||
# All get_state calls must return the correct accumulated value
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 10 # 5 human + 5 AI
|
||||
|
||||
|
||||
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
|
||||
from typing import Annotated
|
||||
@@ -435,16 +322,10 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dict-reducer tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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)])
|
||||
@@ -457,12 +338,14 @@ def test_delta_channel_dict_reducer_fresh_channel() -> None:
|
||||
return {**left, **right}
|
||||
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
# Should be available (not raise EmptyChannelError) and start empty
|
||||
assert ch.is_available()
|
||||
assert ch.get() == {}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_basic_updates() -> None:
|
||||
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
@@ -499,7 +382,7 @@ def test_delta_channel_dict_reducer_writes_reconstruction() -> None:
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_deletions() -> None:
|
||||
"""Dict reducer that treats None values as deletions works end-to-end."""
|
||||
"""Dict reducer that treats None values as deletions works end-to-end (deepagents pattern)."""
|
||||
|
||||
def merge_files(left: dict | None, right: dict) -> dict:
|
||||
if left is None:
|
||||
@@ -513,10 +396,15 @@ def test_delta_channel_dict_reducer_with_deletions() -> None:
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING)
|
||||
|
||||
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
|
||||
|
||||
# Delete file1, add file3
|
||||
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
||||
|
||||
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
# Confirm writes reconstruction produces the same result
|
||||
spec = _delta_channel_with_type(merge_files, dict)
|
||||
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch2.replay_writes(
|
||||
@@ -538,6 +426,7 @@ def test_delta_channel_dict_reducer_overwrite_in_update() -> None:
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
ch.update([{"a": 1}])
|
||||
ch.update([Overwrite({"b": 2, "c": 3})])
|
||||
|
||||
assert ch.get() == {"b": 2, "c": 3}
|
||||
|
||||
|
||||
@@ -561,12 +450,16 @@ def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
|
||||
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`."""
|
||||
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`.
|
||||
|
||||
This is the shape the deepagents filesystem middleware uses for its
|
||||
`files` field; without unwrapping NotRequired we'd fall through to `list`
|
||||
and blow up on the first dict operator call.
|
||||
"""
|
||||
from typing import Annotated
|
||||
|
||||
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:
|
||||
@@ -574,7 +467,10 @@ def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
|
||||
return dict(right)
|
||||
return {**left, **right}
|
||||
|
||||
annotation = Annotated[NotRequired[dict[str, int]], DeltaChannel(merge_dicts)]
|
||||
annotation = Annotated[
|
||||
NotRequired[dict[str, int]],
|
||||
DeltaChannel(merge_dicts),
|
||||
]
|
||||
ch = _get_channel("files", annotation).from_checkpoint(MISSING)
|
||||
assert ch.get() == {}
|
||||
ch.update([{"a": 1}])
|
||||
@@ -583,13 +479,16 @@ def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
|
||||
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel."""
|
||||
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel.
|
||||
|
||||
Mirrors the deepagents filesystem pattern: `files: Annotated[dict, reducer]`
|
||||
where the reducer merges dicts and treats None values as deletions.
|
||||
"""
|
||||
from typing import Annotated
|
||||
|
||||
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:
|
||||
@@ -623,9 +522,12 @@ def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
|
||||
for _ in range(3):
|
||||
graph.invoke({"files": {}}, config)
|
||||
|
||||
# Checkpoint stores only the sentinel — per-step writes live in checkpoint_writes.
|
||||
saved = saver.get_tuple(config)
|
||||
assert saved is not None
|
||||
assert saved.checkpoint["channel_values"]["files"] is DELTA_SENTINEL
|
||||
cv = saved.checkpoint["channel_values"]["files"]
|
||||
assert cv is DELTA_SENTINEL
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert state.values["files"] == {
|
||||
"/doc_1.txt": "content for turn 1",
|
||||
@@ -633,6 +535,7 @@ def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
|
||||
"/doc_3.txt": "content for turn 3",
|
||||
}
|
||||
|
||||
# Deletion path must round-trip through writes replay.
|
||||
def delete_file(state: State) -> dict:
|
||||
return {"files": {"/doc_1.txt": None}}
|
||||
|
||||
|
||||
@@ -1,22 +1,22 @@
|
||||
"""Benchmark: DeltaChannel snapshot_frequency — storage vs. read-depth tradeoff.
|
||||
"""Benchmark: DeltaChannel vs BinaryOperatorAggregate storage and time.
|
||||
|
||||
Run directly: python tests/test_delta_channel_benchmark.py
|
||||
Run via pytest: pytest tests/test_delta_channel_benchmark.py -s
|
||||
|
||||
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.
|
||||
Simulates realistic multi-turn conversations with paragraph-length messages
|
||||
(~100 tokens each) scaling up to 1M-token-equivalent histories.
|
||||
|
||||
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)
|
||||
Token estimates: 1 token ≈ 4 chars; each turn ≈ 200 tokens (human + AI).
|
||||
A 1M-token conversation ≈ 5,000 turns of realistic messages.
|
||||
|
||||
DeltaChannel stores only a zero-byte sentinel in checkpoint_blobs; the actual
|
||||
write data lives in checkpoint_writes (already stored there). Reconstruction
|
||||
walks the parent chain and replays writes through the operator — O(N) total
|
||||
storage vs O(N²) for plain add_messages.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
@@ -26,15 +26,29 @@ 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
|
||||
|
||||
_SQLITE_AVAILABLE = True
|
||||
except ImportError:
|
||||
_SQLITE_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
_POSTGRES_AVAILABLE = True
|
||||
_POSTGRES_URI = "postgres://sydney_runkle@localhost:5441/postgres?sslmode=disable"
|
||||
_POSTGRES_URI = (
|
||||
"postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
|
||||
)
|
||||
except ImportError:
|
||||
_POSTGRES_AVAILABLE = False
|
||||
|
||||
@@ -117,18 +131,6 @@ 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
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -172,8 +174,9 @@ def _run_turns(
|
||||
"""Run n_turns conversation turns.
|
||||
|
||||
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
|
||||
Read latency is the average of 5 get_state calls after the full history
|
||||
is built — forces state rehydration including ancestor replay if needed.
|
||||
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.
|
||||
"""
|
||||
graph = _make_graph(state_cls, checkpointer)
|
||||
config = {"configurable": {"thread_id": "bench"}}
|
||||
@@ -186,16 +189,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
|
||||
|
||||
blob_bytes = (
|
||||
_total_blob_bytes(graph.checkpointer)
|
||||
if isinstance(graph.checkpointer, MemorySaver)
|
||||
else -1
|
||||
)
|
||||
if isinstance(graph.checkpointer, MemorySaver):
|
||||
blob_bytes = _total_blob_bytes(graph.checkpointer)
|
||||
else:
|
||||
blob_bytes = -1
|
||||
return write_elapsed, read_elapsed, blob_bytes
|
||||
|
||||
|
||||
@@ -208,6 +211,7 @@ 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"
|
||||
@@ -217,56 +221,74 @@ def _approx_tokens(n_turns: int) -> str:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpointer factories
|
||||
# Benchmark matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# 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]
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _pg_saver(thread_id: str = "bench"):
|
||||
"""Context manager that yields a fresh PostgresSaver and cleans up after."""
|
||||
with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver:
|
||||
saver.setup()
|
||||
with saver._cursor() as cur:
|
||||
for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
|
||||
cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,))
|
||||
yield saver
|
||||
with saver._cursor() as cur:
|
||||
for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
|
||||
cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,))
|
||||
# Deep-thread counts where add_messages blob storage would exceed 1 GB;
|
||||
# only DeltaChannel runs here.
|
||||
DELTA_ONLY_TURN_COUNTS = [1000]
|
||||
|
||||
|
||||
def _checkpointers() -> list[tuple[str, Any]]:
|
||||
"""Return (label, saver_or_None) pairs for available checkpointers."""
|
||||
result: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
def _checkpointer_factories() -> list[tuple[str, Any]]:
|
||||
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
|
||||
return [("InMemory", None)]
|
||||
|
||||
|
||||
def run_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()
|
||||
result.append(("Postgres", "postgres"))
|
||||
checkpointers.append(("Postgres (plain SELECT)", "postgres"))
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
|
||||
for cp_label, cp_hint in checkpointers:
|
||||
print(f"--- Checkpointer: {cp_label} ---")
|
||||
_run_benchmark_for_checkpointer(cp_hint)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part 1: baseline DeltaChannel(inf) vs add_messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
DELTA_ONLY_TURN_COUNTS = [1000]
|
||||
|
||||
|
||||
def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
|
||||
W = 72
|
||||
def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
import contextlib
|
||||
import tempfile
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
return contextlib.nullcontext(None)
|
||||
return _pg_saver()
|
||||
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
|
||||
|
||||
rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = []
|
||||
for turns in BASELINE_TURN_COUNTS:
|
||||
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:
|
||||
@@ -277,19 +299,26 @@ def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
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)
|
||||
|
||||
def _bytes_or_na(v: Any) -> str:
|
||||
if v is None or v < 0:
|
||||
if v is None:
|
||||
return "n/a"
|
||||
if 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(f"\n [{cp_label}] Storage (blob bytes)")
|
||||
print(
|
||||
f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}"
|
||||
)
|
||||
print(" " + "-" * (W - 2))
|
||||
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"
|
||||
@@ -297,132 +326,65 @@ def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
|
||||
ratio = b_bytes / d_bytes if d_bytes else float("inf")
|
||||
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} {ratio_str:>8}"
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} "
|
||||
f"{ratio_str:>8}"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
print(f"\n [{cp_label}] Read latency (avg of 5 get_state calls)")
|
||||
print(f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12}")
|
||||
print(" " + "-" * (W - 2))
|
||||
# ── 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)
|
||||
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"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
|
||||
)
|
||||
|
||||
|
||||
def run_baseline_benchmark() -> None:
|
||||
print()
|
||||
print("Part 1 — DeltaChannel(inf) vs add_messages: storage & latency")
|
||||
print("=" * 72)
|
||||
for cp_label, cp_hint in _checkpointers():
|
||||
_run_baseline_for_checkpointer(cp_label, cp_hint)
|
||||
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:
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part 2: snapshot_frequency sweep
|
||||
# ---------------------------------------------------------------------------
|
||||
def _per(wt: Any) -> str:
|
||||
if wt is None:
|
||||
return "n/a"
|
||||
return f"{(wt / turns) * 1000:.1f}ms"
|
||||
|
||||
# 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_sweep_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
return contextlib.nullcontext(None)
|
||||
return _pg_saver()
|
||||
|
||||
# Collect results: {turns: {freq_label: (write_s, read_s, bytes)}}
|
||||
results: dict[int, dict[str, tuple[float, float, int]]] = {}
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
results[turns] = {}
|
||||
for freq in SNAPSHOT_FREQUENCIES:
|
||||
state_cls = _make_delta_state(freq)
|
||||
with _make_saver() as saver:
|
||||
wt, rt, bb = _run_turns(turns, state_cls, saver)
|
||||
results[turns][_freq_label(freq)] = (wt, rt, bb)
|
||||
|
||||
freq_labels = [_freq_label(f) for f in SNAPSHOT_FREQUENCIES]
|
||||
col_w = 12
|
||||
|
||||
header = f" {'turns':>6} {'ctx':>10}" + "".join(
|
||||
f" {f'freq={freq_label}':>{col_w}}" for freq_label in freq_labels
|
||||
)
|
||||
|
||||
print(f"\n [{cp_label}] Storage (blob bytes) — lower is better")
|
||||
print(header)
|
||||
print(" " + "-" * (len(header) - 2))
|
||||
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(f"\n [{cp_label}] Read latency (avg of 5 get_state) — lower is better")
|
||||
print(header)
|
||||
print(" " + "-" * (len(header) - 2))
|
||||
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(
|
||||
f"\n [{cp_label}] Per-invoke write latency (total / turns) — lower is better"
|
||||
)
|
||||
print(header)
|
||||
print(" " + "-" * (len(header) - 2))
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
for cp_label, cp_hint in _checkpointers():
|
||||
_run_sweep_for_checkpointer(cp_label, cp_hint)
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_per(b_wt):>12} {_per(d_wt):>12}"
|
||||
)
|
||||
print("=" * W)
|
||||
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(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
|
||||
print(" delta = Annotated[list, DeltaChannel(add_messages)] — O(N) storage")
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry points
|
||||
# Pytest entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
|
||||
)
|
||||
def test_delta_channel_baseline_benchmark(capsys: Any) -> None:
|
||||
"""DeltaChannel(inf) uses less storage than add_messages at scale."""
|
||||
def test_delta_channel_benchmark(capsys: Any) -> None:
|
||||
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
|
||||
with capsys.disabled():
|
||||
run_baseline_benchmark()
|
||||
run_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)
|
||||
@@ -432,41 +394,10 @@ def test_delta_channel_baseline_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_baseline_benchmark()
|
||||
run_snapshot_freq_benchmark()
|
||||
run_benchmark()
|
||||
sys.exit(0)
|
||||
|
||||
@@ -50,9 +50,14 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
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
|
||||
|
||||
|
||||
|
||||
@@ -41,7 +41,6 @@ from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
|
||||
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
|
||||
from langgraph.channels.topic import Topic
|
||||
@@ -80,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__)
|
||||
@@ -9584,90 +9591,3 @@ async def test_delta_channel_update_by_id_end_to_end() -> None:
|
||||
assert "h1" in ids # h1 persists (updated, not duplicated)
|
||||
assert "h2" in ids
|
||||
assert ids.count("h1") == 1, "h1 must not be duplicated"
|
||||
|
||||
|
||||
async def test_delta_channel_write_flushed_before_put() -> None:
|
||||
"""checkpoint_writes are flushed synchronously before put when DELTA_SENTINEL
|
||||
is present, ensuring writes are durable before the sentinel blob is committed.
|
||||
|
||||
We verify this by intercepting put_writes and put calls and confirming
|
||||
put_writes always completes before put is called for sentinel checkpoints.
|
||||
"""
|
||||
import threading
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
i = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"r{i}", id=f"ai{i}")]}
|
||||
|
||||
order: list[str] = []
|
||||
lock = threading.Lock()
|
||||
original_put_writes = InMemorySaver.put_writes
|
||||
original_put = InMemorySaver.put
|
||||
|
||||
def tracked_put_writes(self, config, writes, task_id, task_path=""):
|
||||
result = original_put_writes(self, config, writes, task_id, task_path)
|
||||
with lock:
|
||||
order.append("put_writes")
|
||||
return result
|
||||
|
||||
def tracked_put(self, config, checkpoint, metadata, new_versions):
|
||||
# Check if this checkpoint has any DELTA_SENTINEL blobs
|
||||
has_sentinel = any(
|
||||
v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values()
|
||||
)
|
||||
if has_sentinel:
|
||||
with lock:
|
||||
order.append("put_sentinel")
|
||||
else:
|
||||
with lock:
|
||||
order.append("put_snapshot")
|
||||
return original_put(self, config, checkpoint, metadata, new_versions)
|
||||
|
||||
InMemorySaver.put_writes = tracked_put_writes
|
||||
InMemorySaver.put = tracked_put
|
||||
try:
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "flush-test"}}
|
||||
|
||||
for i in range(3):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config
|
||||
)
|
||||
|
||||
# For every sentinel put, all preceding put_writes must already be in order
|
||||
for i, event in enumerate(order):
|
||||
if event == "put_sentinel":
|
||||
# All put_writes before this index must appear before this sentinel
|
||||
preceding = order[:i]
|
||||
assert "put_writes" in preceding, (
|
||||
f"put_sentinel at index {i} had no preceding put_writes: {order}"
|
||||
)
|
||||
# And the most recent put_writes must come before this sentinel
|
||||
last_write_idx = max(
|
||||
j for j, e in enumerate(order[:i]) if e == "put_writes"
|
||||
)
|
||||
assert last_write_idx < i, (
|
||||
f"put_writes at {last_write_idx} not before put_sentinel at {i}"
|
||||
)
|
||||
finally:
|
||||
InMemorySaver.put_writes = original_put_writes
|
||||
InMemorySaver.put = original_put
|
||||
|
||||
# Final state must still be correct
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 6 # 3 human + 3 AI
|
||||
|
||||
Reference in New Issue
Block a user