mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 11:19:54 +02:00
Compare commits
40
Commits
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,405 @@
|
||||
# DiffChannel: Incremental Checkpoint Storage for Append-Style Reducers
|
||||
|
||||
**Date:** 2026-04-17
|
||||
**Status:** Approved for implementation
|
||||
**Scope:** `libs/checkpoint`, `libs/langgraph`, `libs/checkpoint-postgres`
|
||||
|
||||
---
|
||||
|
||||
## Motivation
|
||||
|
||||
LangGraph checkpoints today store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, this means each checkpoint blob contains the entire conversation history. Storage cost grows O(N²) in the number of turns: step 1 stores 1 message, step 100 stores 100 messages, step 1000 stores 1000 messages. For long-running agentic conversations with high-token messages this is untenable.
|
||||
|
||||
The fix is to store only the **delta** (new writes) per step, reconstructing the full accumulated value at load time by replaying the chain. This is an opt-in mechanism — existing graphs are unaffected.
|
||||
|
||||
---
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- **Compaction / materialized snapshots**: deferred. Load cost stays O(N) blob fetches but those fetches are batched into a single query — acceptable for now.
|
||||
- **SQLite saver support**: SQLite stores all channel values inline in one row (no per-channel blob table). Deferred to a follow-up.
|
||||
- **Automatic migration** of existing `BinaryOperatorAggregate` channels: users opt in explicitly. Old checkpoints load correctly via the backwards-compatibility path in `from_checkpoint`.
|
||||
|
||||
---
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
```
|
||||
User state definition
|
||||
└── Annotated[list[AnyMessage], DiffChannel(add_messages)]
|
||||
|
||||
Write path (per superstep)
|
||||
DiffChannel.update() — apply operator, accumulate writes in _pending
|
||||
DiffChannel.checkpoint() — return DiffDelta(delta=_pending, prev_version=_base_version)
|
||||
serde.dumps_typed() — serialize DiffDelta as ("diff", msgpack_bytes)
|
||||
saver.put() — store blob at (thread_id, ns, "messages", version_N)
|
||||
DiffChannel.after_checkpoint(version_N) — advance _base_version, clear _pending
|
||||
|
||||
Read path (on graph load or time-travel)
|
||||
saver.get_tuple() — fetch current-version blob per channel
|
||||
saver._load_blobs() — detect "diff" type → follow chain to reconstruct DiffChainValue
|
||||
DiffChannel.from_checkpoint(DiffChainValue) — replay deltas with operator → full list
|
||||
DiffChannel.after_checkpoint(version_N) — set _base_version for next write
|
||||
```
|
||||
|
||||
The pregel layer (`_checkpoint.py`, `_loop.py`) is unchanged except for two small additions to call the new `after_checkpoint` hook. The saver public interface (`BaseCheckpointSaver`) gains no new methods. All chain-following logic lives inside each saver's private `_load_blobs`.
|
||||
|
||||
---
|
||||
|
||||
## New Protocol Types
|
||||
|
||||
**Location:** `libs/checkpoint/langgraph/checkpoint/base/__init__.py`
|
||||
|
||||
Two dataclasses form the contract between `DiffChannel` and savers:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class DiffDelta:
|
||||
"""Returned by DiffChannel.checkpoint(). Written to the blob store."""
|
||||
delta: list[Any] # raw writes passed to update() this step
|
||||
prev_version: str | None # version of the previous diff blob; None = chain root
|
||||
```
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class DiffChainValue:
|
||||
"""Passed to DiffChannel.from_checkpoint(). Assembled by _load_blobs()."""
|
||||
base: list[Any] | None # starting accumulated value (None = empty start)
|
||||
deltas: list[list[Any]] # write-sets ordered oldest → newest
|
||||
```
|
||||
|
||||
`DiffDelta` lives in the checkpoint base package (not the channel module) so savers can import it without creating a circular dependency. `DiffChainValue` is there for the same reason.
|
||||
|
||||
---
|
||||
|
||||
## `BaseChannel.after_checkpoint()` Hook
|
||||
|
||||
**Location:** `libs/langgraph/langgraph/channels/base.py`
|
||||
|
||||
```python
|
||||
def after_checkpoint(self, version: Any) -> None:
|
||||
"""Called after checkpoint() (with the new version) and after from_checkpoint()
|
||||
(with the current version). No-op by default; DiffChannel overrides."""
|
||||
pass
|
||||
```
|
||||
|
||||
This is a **non-abstract, no-op default** — fully backwards compatible. All existing channels inherit it silently. It is NOT in the abstract interface.
|
||||
|
||||
---
|
||||
|
||||
## `DiffChannel[V]`
|
||||
|
||||
**Location:** `libs/langgraph/langgraph/channels/diff.py` (new file)
|
||||
|
||||
### Internal state
|
||||
|
||||
| Attribute | Type | Description |
|
||||
|---|---|---|
|
||||
| `value` | `list[V]` | Full accumulated value (the reconstructed list) |
|
||||
| `operator` | `Callable` | The binary reducer (e.g. `add_messages`) |
|
||||
| `_pending` | `list[Any]` | Raw writes accumulated since last `after_checkpoint` call |
|
||||
| `_base_version` | `str \| None` | Version this channel was last checkpointed at (= `prev_version` for next delta) |
|
||||
| `_overwritten` | `bool` | True if an `Overwrite` was applied since last `after_checkpoint`; makes next blob a chain root |
|
||||
|
||||
### `update(values)`
|
||||
|
||||
Mirrors `BinaryOperatorAggregate.update()` with two additions:
|
||||
|
||||
1. For each non-Overwrite value: apply `self.operator(self.value, value)` as before; **also append the raw incoming value to `self._pending`**.
|
||||
2. For an `Overwrite(v)` value: set `self.value = v`; set `self._pending = list(v)` (full value becomes the new delta); set `self._overwritten = True`.
|
||||
|
||||
The key: `_pending` stores the **incoming writes** (what was passed to `update()`), not the diff of `self.value`. This is important because `add_messages` handles removal and update-by-ID — replaying the writes with `operator` during reconstruction applies that logic correctly.
|
||||
|
||||
### `checkpoint()`
|
||||
|
||||
```python
|
||||
def checkpoint(self) -> DiffDelta:
|
||||
return DiffDelta(
|
||||
delta=self._pending[:],
|
||||
prev_version=None if self._overwritten else self._base_version,
|
||||
)
|
||||
```
|
||||
|
||||
- Normal step: `prev_version = self._base_version` → chain link
|
||||
- After Overwrite: `prev_version = None` → chain root (reconstruction stops here and uses `delta` as the full base value)
|
||||
|
||||
Returns `DiffDelta`, never the raw accumulated list. The serde handles serialization.
|
||||
|
||||
### `from_checkpoint(checkpoint)`
|
||||
|
||||
```python
|
||||
def from_checkpoint(self, checkpoint) -> Self:
|
||||
new = DiffChannel(self.typ, self.operator)
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
new.value = []
|
||||
elif isinstance(checkpoint, DiffChainValue):
|
||||
accumulated = checkpoint.base or []
|
||||
for step_writes in checkpoint.deltas:
|
||||
# Mirror update() exactly: apply each write individually so operator
|
||||
# semantics (e.g. add_messages ID-based removal) are respected.
|
||||
for write in step_writes:
|
||||
accumulated = new.operator(accumulated, write)
|
||||
new.value = accumulated
|
||||
elif isinstance(checkpoint, DiffDelta):
|
||||
# Unsupported saver: _load_blobs returned a raw DiffDelta instead of
|
||||
# assembling a DiffChainValue. Raise rather than silently losing history.
|
||||
raise ValueError(
|
||||
"DiffChannel received a raw DiffDelta from the checkpoint saver. "
|
||||
"Your saver does not support incremental channel storage. "
|
||||
"Use InMemorySaver or PostgresSaver."
|
||||
)
|
||||
else:
|
||||
# Backwards compat: plain list from old BinaryOperatorAggregate checkpoint.
|
||||
new.value = checkpoint
|
||||
new._pending = []
|
||||
new._base_version = None # set by the subsequent after_checkpoint() call
|
||||
return new
|
||||
```
|
||||
|
||||
The operator is available on `self` (the channel spec) so reconstruction is correct for any reducer — the saver never needs to know about `add_messages`.
|
||||
|
||||
`_pending` stores **individual writes** (each `value` from `update()`'s `values` sequence), so each `step_writes` list in `DiffChainValue.deltas` is replayed write-by-write — identical to the `update()` loop.
|
||||
|
||||
### `after_checkpoint(version)`
|
||||
|
||||
```python
|
||||
def after_checkpoint(self, version: Any) -> None:
|
||||
if version != self._base_version:
|
||||
self._base_version = version
|
||||
self._pending = []
|
||||
self._overwritten = False
|
||||
```
|
||||
|
||||
No-op when `version == self._base_version` (channel wasn't updated this step — blob was not written). Clears `_pending` and advances `_base_version` when the channel was actually checkpointed.
|
||||
|
||||
### Opt-in API
|
||||
|
||||
```python
|
||||
from langgraph.channels.diff import DiffChannel
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DiffChannel(add_messages)]
|
||||
```
|
||||
|
||||
`StateGraph` already handles `BaseChannel` instances as annotation metadata — `DiffChannel` inherits this without any changes to `StateGraph`.
|
||||
|
||||
---
|
||||
|
||||
## Serde Extension
|
||||
|
||||
**Location:** `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py`
|
||||
|
||||
Add one branch to `dumps_typed` (before the `else` msgpack fallback), using the existing module-level `_msgpack_enc` so message ext-types (Pydantic v2, etc.) are handled correctly:
|
||||
|
||||
```python
|
||||
elif isinstance(obj, DiffDelta):
|
||||
return "diff", _msgpack_enc({"d": obj.delta, "p": obj.prev_version})
|
||||
```
|
||||
|
||||
Add one branch to `loads_typed` so savers can decode diff blobs without importing `ormsgpack` directly:
|
||||
|
||||
```python
|
||||
elif type_ == "diff":
|
||||
return ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# returns {"d": [writes...], "p": prev_version_str_or_none}
|
||||
```
|
||||
|
||||
Savers call `serde.loads_typed(("diff", raw_bytes))` to decode a diff blob into `{"d": ..., "p": ...}`, then check `type_tag == "diff"` to trigger chain traversal. The serde layer is the only place that knows about `ormsgpack`.
|
||||
|
||||
---
|
||||
|
||||
## Saver Changes
|
||||
|
||||
### InMemorySaver
|
||||
|
||||
**`put()` — `libs/checkpoint/langgraph/checkpoint/memory/__init__.py`**
|
||||
|
||||
No change needed. The existing `self.serde.dumps_typed(values[k])` call already handles `DiffDelta` via the new serde branch above, storing it as `("diff", bytes)`.
|
||||
|
||||
**`_load_blobs()` — same file**
|
||||
|
||||
After checking `vv[0] != "empty"`, add a branch for `"diff"` before calling `serde.loads_typed`:
|
||||
|
||||
```python
|
||||
def _load_blobs(self, thread_id, checkpoint_ns, versions):
|
||||
channel_values = {}
|
||||
diff_channels = {} # channel_name -> current_version for diff channels
|
||||
|
||||
for k, v in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, v)
|
||||
if kk not in self.blobs:
|
||||
continue
|
||||
type_tag, blob_bytes = self.blobs[kk]
|
||||
if type_tag == "diff":
|
||||
diff_channels[k] = v # handle below
|
||||
elif type_tag != "empty":
|
||||
channel_values[k] = self.serde.loads_typed((type_tag, blob_bytes))
|
||||
|
||||
for k, current_version in diff_channels.items():
|
||||
# Follow chain: newest → oldest, then reverse
|
||||
chain_deltas = []
|
||||
base = None
|
||||
version = current_version
|
||||
while version is not None:
|
||||
kk = (thread_id, checkpoint_ns, k, version)
|
||||
if kk not in self.blobs:
|
||||
break
|
||||
type_tag, blob_bytes = self.blobs[kk]
|
||||
if type_tag == "diff":
|
||||
# Use serde so we don't need to import ormsgpack directly
|
||||
payload = self.serde.loads_typed((type_tag, blob_bytes))
|
||||
chain_deltas.append(payload["d"])
|
||||
version = payload["p"] # prev_version; None = root
|
||||
else:
|
||||
# Old non-diff blob encountered: treat as base accumulated value
|
||||
base = self.serde.loads_typed((type_tag, blob_bytes))
|
||||
break
|
||||
chain_deltas.reverse()
|
||||
channel_values[k] = DiffChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
return channel_values
|
||||
```
|
||||
|
||||
Each blob lookup is O(1) on the dict. Total: N dict lookups for a chain of depth N. Memory usage is identical to loading a single full-list blob (same total bytes, split across N entries).
|
||||
|
||||
### PostgresSaver
|
||||
|
||||
**`_load_blobs()` — `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py`**
|
||||
|
||||
The existing `SELECT_SQL` fetches one blob per channel via a JOIN. After running that query, detect any `"diff"` channels in the result and issue one additional range query:
|
||||
|
||||
```python
|
||||
def _load_blobs(self, blob_values):
|
||||
if not blob_values:
|
||||
return {}
|
||||
|
||||
result = {}
|
||||
diff_channels = {} # channel_name -> current_version (as str)
|
||||
|
||||
for k, t, v in blob_values:
|
||||
channel = k.decode()
|
||||
type_tag = t.decode()
|
||||
if type_tag == "diff":
|
||||
# Decode via serde — no direct ormsgpack import needed
|
||||
payload = self.serde.loads_typed((type_tag, v))
|
||||
diff_channels[channel] = payload # store for chain fetch
|
||||
elif type_tag != "empty":
|
||||
result[channel] = self.serde.loads_typed((type_tag, v))
|
||||
|
||||
if diff_channels:
|
||||
result.update(self._load_diff_chains(diff_channels))
|
||||
|
||||
return result
|
||||
```
|
||||
|
||||
`_load_diff_chains` issues one SQL query per diff channel (typically just `messages`):
|
||||
|
||||
```sql
|
||||
SELECT version, type, blob
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = %s
|
||||
AND checkpoint_ns = %s
|
||||
AND channel = %s
|
||||
AND version <= %s
|
||||
ORDER BY version ASC
|
||||
```
|
||||
|
||||
In Python, iterate rows in ascending version order: if `type = "diff"`, accumulate the delta; if any other type is encountered, treat it as the base accumulated value and stop. Return `DiffChainValue(base=..., deltas=[...])`.
|
||||
|
||||
This results in **at most 2 queries total** for a graph with one `DiffChannel` — existing behaviour for all other channels is unchanged.
|
||||
|
||||
**`put()` / `_dump_blobs()`**
|
||||
|
||||
No change needed. `_dump_blobs` calls `self.serde.dumps_typed(v)` for each channel value in `new_versions`. When `v` is a `DiffDelta`, the serde produces `("diff", bytes)` which is stored as `type = "diff"` in `checkpoint_blobs`. The `ON CONFLICT DO NOTHING` semantics are preserved.
|
||||
|
||||
### SQLite
|
||||
|
||||
Deferred. `SqliteSaver` stores the entire checkpoint as a single serialized row — it has no per-channel blob table. Supporting `DiffChannel` on SQLite would require adding a new blobs table, which is a separate migration tracked separately.
|
||||
|
||||
---
|
||||
|
||||
## Pregel Layer Changes
|
||||
|
||||
### `channels_from_checkpoint` — `libs/langgraph/langgraph/pregel/_checkpoint.py`
|
||||
|
||||
After constructing each channel from its checkpoint value, call `after_checkpoint` so the channel records its current version:
|
||||
|
||||
```python
|
||||
channels = {}
|
||||
for k, v in channel_specs.items():
|
||||
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
ch.after_checkpoint(checkpoint["channel_versions"].get(k))
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
```
|
||||
|
||||
Existing channels get the no-op `after_checkpoint`. `DiffChannel` uses it to set `_base_version`.
|
||||
|
||||
### `PregelLoop._put_checkpoint` — `libs/langgraph/langgraph/pregel/_loop.py`
|
||||
|
||||
After `create_checkpoint(self.checkpoint, self.channels, self.step, ...)` returns and `do_checkpoint is True` and `self.channels is not None`, iterate channels and notify:
|
||||
|
||||
```python
|
||||
if do_checkpoint and self.channels:
|
||||
for k, ch in self.channels.items():
|
||||
ch.after_checkpoint(self.checkpoint["channel_versions"].get(k))
|
||||
```
|
||||
|
||||
This is called after `create_checkpoint` updates `self.checkpoint["channel_versions"]`, so `get(k)` returns the new version for updated channels and the old version for unchanged ones. `DiffChannel.after_checkpoint` only clears `_pending` when `version != _base_version`, so unchanged channels are no-ops.
|
||||
|
||||
---
|
||||
|
||||
## Backwards Compatibility
|
||||
|
||||
| Scenario | Behaviour |
|
||||
|---|---|
|
||||
| Existing graph using `add_messages` (BinaryOperatorAggregate) | Unaffected — no code changes, no data migration |
|
||||
| New graph with `DiffChannel`, loading old checkpoint blobs | `from_checkpoint` receives a plain `list` → used directly as accumulated value |
|
||||
| `DiffChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported |
|
||||
| `DiffChannel` with `SqliteSaver` | `from_checkpoint` receives a raw `DiffDelta` (SqliteSaver stores channel_values inline), raises `ValueError` with a clear message pointing to supported savers |
|
||||
| Time-travel / fork to past checkpoint | Chain traversal uses the version at that checkpoint → reconstruction is correct |
|
||||
| `update_state` | Treated as a normal step: writes are deltas chained to history |
|
||||
| `Overwrite` value | Resets chain: next blob has `prev_version=None`; reconstruction starts fresh |
|
||||
|
||||
---
|
||||
|
||||
## Testing Strategy
|
||||
|
||||
1. **Unit tests for `DiffChannel`** (`libs/langgraph/tests/`):
|
||||
- `update` → `checkpoint` → `after_checkpoint` → `checkpoint` lifecycle (2 steps, verify delta isolation)
|
||||
- `from_checkpoint(DiffChainValue)` correctly replays multi-step chains using the operator
|
||||
- `from_checkpoint(plain_list)` backwards-compat path
|
||||
- `Overwrite` creates a root blob (`prev_version=None`) and reconstruction ignores prior chain
|
||||
- `after_checkpoint` no-ops when version is unchanged
|
||||
|
||||
2. **Integration tests with `InMemorySaver`** (`libs/langgraph/tests/`):
|
||||
- 10-step conversation: verify final loaded state equals full accumulated messages
|
||||
- Time-travel: fork to step 5, verify only messages 1–5 are present
|
||||
- Mixed graph: some channels `BinaryOperatorAggregate`, one `DiffChannel` — both reconstruct correctly
|
||||
|
||||
3. **Serde tests** (`libs/checkpoint/tests/`):
|
||||
- `DiffDelta` round-trips through `dumps_typed` / saver storage
|
||||
- Old `"msgpack"` blob for a channel → `DiffChannel.from_checkpoint` handles it
|
||||
|
||||
4. **Postgres integration tests** (`libs/checkpoint-postgres/tests/`):
|
||||
- Range query reconstructs correct full list after N steps
|
||||
- Time-travel to checkpoint M reconstructs correct list of M messages
|
||||
|
||||
---
|
||||
|
||||
## Files Changed
|
||||
|
||||
| File | Change |
|
||||
|---|---|
|
||||
| `libs/checkpoint/langgraph/checkpoint/base/__init__.py` | Add `DiffDelta`, `DiffChainValue` dataclasses |
|
||||
| `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` | Add `"diff"` branch in `dumps_typed` |
|
||||
| `libs/checkpoint/langgraph/checkpoint/memory/__init__.py` | Chain traversal in `_load_blobs` |
|
||||
| `libs/langgraph/langgraph/channels/base.py` | Add no-op `after_checkpoint` method |
|
||||
| `libs/langgraph/langgraph/channels/diff.py` | **New file** — `DiffChannel` implementation |
|
||||
| `libs/langgraph/langgraph/channels/__init__.py` | Export `DiffChannel` |
|
||||
| `libs/langgraph/langgraph/pregel/_checkpoint.py` | Call `after_checkpoint` in `channels_from_checkpoint` |
|
||||
| `libs/langgraph/langgraph/pregel/_loop.py` | Call `after_checkpoint` after `create_checkpoint` |
|
||||
| `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` | Range-query chain reconstruction in `_load_blobs` |
|
||||
@@ -8,13 +8,11 @@ from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -25,12 +23,7 @@ from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_BLOBS_SQL,
|
||||
SELECT_DELTA_PARENTS_SQL,
|
||||
SELECT_DELTA_WRITES_SQL,
|
||||
BasePostgresSaver,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
|
||||
Conn = _internal.Conn # For backward compatibility
|
||||
@@ -437,41 +430,42 @@ class PostgresSaver(BasePostgresSaver):
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def _get_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
|
||||
|
||||
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
|
||||
`checkpoint_blobs`) each filtered by `(thread_id, checkpoint_ns)` and
|
||||
the per-table key. Plain SELECTs let the planner pick straight index
|
||||
scans; rationale + benchmark in `notes/delta_channel_query_bench.md`.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
# Caller didn't specify a target — resolve to the latest
|
||||
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
|
||||
# returns the newest; its config carries the resolved id.
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
def get_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
"""Look up a channel blob by checkpoint ID + channel via checkpoint_blobs."""
|
||||
with self._cursor() as cur:
|
||||
cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns))
|
||||
parents_rows = cur.fetchall()
|
||||
cur.execute(SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel))
|
||||
writes_rows = cur.fetchall()
|
||||
cur.execute(SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel))
|
||||
blobs_rows = cur.fetchall()
|
||||
return self._build_delta_channel_writes_history(
|
||||
channel=channel,
|
||||
target_id=checkpoint_id,
|
||||
parents_rows=parents_rows,
|
||||
writes_rows=writes_rows,
|
||||
blobs_rows=blobs_rows,
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT cb.type, cb.blob
|
||||
FROM checkpoint_blobs cb
|
||||
WHERE cb.thread_id = %s
|
||||
AND cb.checkpoint_ns = %s
|
||||
AND cb.channel = %s
|
||||
AND cb.version = (
|
||||
SELECT checkpoint->'channel_versions'->>%s
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
|
||||
)
|
||||
""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return NotImplemented
|
||||
return self.serde.loads_typed((row["type"], row["blob"]))
|
||||
|
||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
@@ -485,7 +479,13 @@ class PostgresSaver(BasePostgresSaver):
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
channel_values = self._load_blobs(value["channel_values"])
|
||||
with self._cursor() as cur:
|
||||
channel_values = self._load_blobs(
|
||||
value["channel_values"],
|
||||
thread_id=value["thread_id"],
|
||||
checkpoint_ns=value["checkpoint_ns"],
|
||||
cur=cur,
|
||||
)
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
|
||||
@@ -8,13 +8,11 @@ from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -25,12 +23,7 @@ from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_BLOBS_SQL,
|
||||
SELECT_DELTA_PARENTS_SQL,
|
||||
SELECT_DELTA_WRITES_SQL,
|
||||
BasePostgresSaver,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
|
||||
Conn = _ainternal.Conn # For backward compatibility
|
||||
@@ -398,44 +391,42 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _aget_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
|
||||
|
||||
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
|
||||
`checkpoint_blobs`); rows assembled by the shared pure helper on
|
||||
`BasePostgresSaver`. Rationale + benchmark in
|
||||
`notes/delta_channel_query_bench.md`.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
target = await self.aget_tuple(config)
|
||||
if target is None:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
async def aget_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
"""Async look up of a channel blob by checkpoint ID + channel name."""
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns)
|
||||
"""
|
||||
SELECT cb.type, cb.blob
|
||||
FROM checkpoint_blobs cb
|
||||
WHERE cb.thread_id = %s
|
||||
AND cb.checkpoint_ns = %s
|
||||
AND cb.channel = %s
|
||||
AND cb.version = (
|
||||
SELECT checkpoint->'channel_versions'->>%s
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
|
||||
)
|
||||
""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
),
|
||||
)
|
||||
parents_rows = await cur.fetchall()
|
||||
await cur.execute(
|
||||
SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel)
|
||||
)
|
||||
writes_rows = await cur.fetchall()
|
||||
await cur.execute(
|
||||
SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel)
|
||||
)
|
||||
blobs_rows = await cur.fetchall()
|
||||
return self._build_delta_channel_writes_history(
|
||||
channel=channel,
|
||||
target_id=checkpoint_id,
|
||||
parents_rows=parents_rows,
|
||||
writes_rows=writes_rows,
|
||||
blobs_rows=blobs_rows,
|
||||
)
|
||||
row = await cur.fetchone()
|
||||
if row is None:
|
||||
return NotImplemented
|
||||
return self.serde.loads_typed((row["type"], row["blob"]))
|
||||
|
||||
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
@@ -452,6 +443,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
thread_id = value["thread_id"]
|
||||
checkpoint_ns = value["checkpoint_ns"]
|
||||
blob_values = value["channel_values"]
|
||||
|
||||
channel_values: dict[str, Any] = {}
|
||||
if blob_values:
|
||||
channel_values = self._load_blobs(blob_values)
|
||||
|
||||
@@ -8,15 +8,12 @@ from typing import Any, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
PendingWrite,
|
||||
_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
|
||||
@@ -155,30 +152,6 @@ INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
|
||||
"""
|
||||
|
||||
# DeltaChannel reconstruction: three plain indexed SELECTs per channel.
|
||||
# Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE
|
||||
# carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are
|
||||
# 3x-100x faster in the realistic depth range and the Python walk is O(n).
|
||||
SELECT_DELTA_PARENTS_SQL = """
|
||||
SELECT checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
"""
|
||||
|
||||
SELECT_DELTA_WRITES_SQL = """
|
||||
SELECT checkpoint_id, type, blob, task_id, idx
|
||||
FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
|
||||
"""
|
||||
|
||||
SELECT_DELTA_BLOBS_SQL = """
|
||||
SELECT version, type, blob
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
|
||||
"""
|
||||
|
||||
|
||||
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
SELECT_SQL = SELECT_SQL
|
||||
@@ -213,106 +186,22 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
def _load_blobs(
|
||||
self,
|
||||
blob_values: Any,
|
||||
blob_values: list[tuple[bytes, bytes, bytes]],
|
||||
*,
|
||||
thread_id: str = "",
|
||||
checkpoint_ns: str = "",
|
||||
cur: Any = None,
|
||||
) -> dict[str, Any]:
|
||||
if not blob_values:
|
||||
return {}
|
||||
result: dict[str, Any] = {}
|
||||
for k, t, v in blob_values:
|
||||
channel = k.decode()
|
||||
type_tag = t.decode()
|
||||
if type_tag != "empty":
|
||||
result[k.decode()] = self.serde.loads_typed((type_tag, v))
|
||||
result[channel] = self.serde.loads_typed((type_tag, v))
|
||||
return result
|
||||
|
||||
def _build_delta_channel_writes_history(
|
||||
self,
|
||||
*,
|
||||
channel: str,
|
||||
target_id: str,
|
||||
parents_rows: Sequence[Any],
|
||||
writes_rows: Sequence[Any],
|
||||
blobs_rows: Sequence[Any],
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Reconstruct one delta channel's history from rows of the three SELECTs.
|
||||
|
||||
Pure data transform shared by sync (`PostgresSaver`) and async
|
||||
(`AsyncPostgresSaver`); both paths run the queries themselves and
|
||||
feed the rows here.
|
||||
|
||||
Walk is newest → oldest from the target's parent. A non-sentinel
|
||||
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
|
||||
walk and is returned as the seed so replay starts from it.
|
||||
|
||||
Writes stored at `target_id` itself are pending writes for the next
|
||||
step and are excluded — the walk begins at the target's parent.
|
||||
"""
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_of: dict[str, str | None] = {}
|
||||
for r in parents_rows:
|
||||
cid = r["checkpoint_id"]
|
||||
parent_of[cid] = r["parent_checkpoint_id"]
|
||||
ver_of[cid] = r["ver"]
|
||||
|
||||
ancestors: list[str] = []
|
||||
cid = parent_of.get(target_id)
|
||||
while cid is not None:
|
||||
ancestors.append(cid)
|
||||
cid = parent_of.get(cid)
|
||||
if not ancestors:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
ancestor_set = set(ancestors)
|
||||
|
||||
# Group writes by ancestor cid; sort within (task_id DESC, idx DESC)
|
||||
# to match the prior CTE ordering — newest write first per ancestor.
|
||||
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
|
||||
for r in writes_rows:
|
||||
cid = r["checkpoint_id"]
|
||||
if cid not in ancestor_set:
|
||||
continue
|
||||
writes_by_cid.setdefault(cid, []).append(
|
||||
(r["type"], r["blob"], r["task_id"], r["idx"])
|
||||
)
|
||||
for ws in writes_by_cid.values():
|
||||
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
|
||||
|
||||
blob_by_ver: dict[str, tuple[str, bytes]] = {
|
||||
r["version"]: (r["type"], r["blob"]) for r in blobs_rows
|
||||
}
|
||||
|
||||
collected: list[PendingWrite] = [] # newest first; reversed at the end
|
||||
for cid in ancestors:
|
||||
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)
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
thread_id: str,
|
||||
|
||||
@@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values(
|
||||
|
||||
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
||||
|
||||
async def patched_load_checkpoint_tuple(value):
|
||||
def patched_load_checkpoint_tuple(value):
|
||||
value["checkpoint"].pop("channel_values", None)
|
||||
return await load_checkpoint_tuple(value)
|
||||
return load_checkpoint_tuple(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
import asyncio
|
||||
import copy
|
||||
import dataclasses
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
|
||||
from typing import (
|
||||
from typing import ( # noqa: UP035
|
||||
Any,
|
||||
Generic,
|
||||
Literal,
|
||||
@@ -19,31 +20,56 @@ from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
DELTA_SENTINEL as DELTA_SENTINEL,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
ERROR,
|
||||
INTERRUPT,
|
||||
RESUME,
|
||||
SCHEDULED,
|
||||
ChannelProtocol,
|
||||
_DeltaSnapshot,
|
||||
)
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
PendingWrite = tuple[str, str, Any]
|
||||
|
||||
# Task-local guard: ContextVar is copied per asyncio Task, so concurrent
|
||||
# requests on the same event-loop thread do not share this flag. A plain
|
||||
# `threading.local()` would leak across tasks and let one in-flight
|
||||
# reconstruction silently short-circuit another.
|
||||
_DELTA_RECONSTRUCTION: contextvars.ContextVar[bool] = contextvars.ContextVar(
|
||||
"_DELTA_RECONSTRUCTION", default=False
|
||||
)
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DeltaValue:
|
||||
"""Returned by DeltaChannel.checkpoint(). Represents one step's writes."""
|
||||
|
||||
delta: list[Any]
|
||||
prev_checkpoint_id: (
|
||||
str | None
|
||||
) # ID of checkpoint containing previous blob; None = chain root
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DeltaChainValue:
|
||||
"""Passed to DeltaChannel.from_checkpoint(). Assembled during checkpoint hydration."""
|
||||
|
||||
base: list[Any] | None # starting accumulated value; None = start from empty
|
||||
deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest
|
||||
|
||||
|
||||
CheckpointHydrationKind = Literal["delta"]
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class IncrementalChannelSpec:
|
||||
"""Describes a checkpoint field that needs saver-side materialization."""
|
||||
|
||||
name: str
|
||||
kind: CheckpointHydrationKind
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class CheckpointHydrationPlan:
|
||||
"""Lists the checkpoint fields eligible for saver-side materialization."""
|
||||
|
||||
channels: tuple[IncrementalChannelSpec, ...]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_MISSING_SENTINEL = object()
|
||||
|
||||
|
||||
# Marked as total=False to allow for future expansion.
|
||||
@@ -134,30 +160,6 @@ class CheckpointTuple(NamedTuple):
|
||||
pending_writes: list[PendingWrite] | None = None
|
||||
|
||||
|
||||
class _ChannelWritesHistory(NamedTuple):
|
||||
"""Result of `BaseCheckpointSaver._get_channel_writes_history`.
|
||||
|
||||
Storage-level view of what one channel wrote across the ancestor chain
|
||||
of a target checkpoint:
|
||||
|
||||
* `seed` — the nearest ancestor's stored blob value for this channel,
|
||||
or `DELTA_SENTINEL` if the walk reached the root without finding a
|
||||
stored value. A non-sentinel seed typically indicates a pre-delta
|
||||
snapshot preserved across a channel-type migration (e.g.
|
||||
`BinaryOperatorAggregate` storage extended under `DeltaChannel`).
|
||||
* `writes` — on-path deltas oldest→newest, one `PendingWrite` per
|
||||
step that wrote to this channel. Writes stored at the target
|
||||
checkpoint itself are pending for the next super-step and are
|
||||
excluded.
|
||||
|
||||
Experimental: method surface may change; the NamedTuple shape is the
|
||||
contract.
|
||||
"""
|
||||
|
||||
seed: Any
|
||||
writes: list[PendingWrite]
|
||||
|
||||
|
||||
class BaseCheckpointSaver(Generic[V]):
|
||||
"""Base class for creating a graph checkpointer.
|
||||
|
||||
@@ -496,120 +498,171 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _get_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""**Experimental.** Query one channel's writes along the parent chain.
|
||||
def materialize_checkpoint(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Materialize any saver-managed incremental values in a checkpoint."""
|
||||
if plan is None or not plan.channels:
|
||||
return checkpoint
|
||||
|
||||
Storage-level query, not channel semantics: returns `(seed, writes)`
|
||||
reflecting what storage knows about a single channel across the
|
||||
ancestor chain of the target checkpoint identified by `config`.
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
current_checkpoint_id = checkpoint.get("id")
|
||||
assembled: dict[str, Any] = {}
|
||||
|
||||
* `writes` — on-path deltas oldest→newest as `PendingWrite` tuples.
|
||||
Writes stored at the target `checkpoint_id` itself are pending
|
||||
for the next super-step and are excluded.
|
||||
* `seed` — the nearest ancestor's stored blob value for this
|
||||
channel; `DELTA_SENTINEL` if the walk reached the root without
|
||||
finding a stored value. A non-sentinel seed typically indicates
|
||||
a pre-delta snapshot preserved across a channel-type migration.
|
||||
for spec in plan.channels:
|
||||
value = checkpoint["channel_values"].get(spec.name)
|
||||
if spec.kind != "delta" or not isinstance(value, DeltaValue):
|
||||
continue
|
||||
|
||||
Walks the **parent chain** (not `list(before=...)`): for forked
|
||||
threads, only on-path ancestors contribute.
|
||||
assembled_value = self._materialize_delta_channel(
|
||||
thread_id=thread_id,
|
||||
checkpoint_ns=checkpoint_ns,
|
||||
current_checkpoint_id=current_checkpoint_id,
|
||||
channel=spec.name,
|
||||
value=value,
|
||||
)
|
||||
if assembled_value is not None:
|
||||
assembled[spec.name] = assembled_value
|
||||
|
||||
Reference implementation walks `get_tuple` + `parent_config`,
|
||||
inspecting each ancestor's `channel_values[channel]` for the seed
|
||||
terminator. Savers with direct storage access (`InMemorySaver`,
|
||||
`PostgresSaver`) override for performance; the return contract is
|
||||
fixed here.
|
||||
if not assembled:
|
||||
return checkpoint
|
||||
|
||||
Underscore-prefixed because the method surface is experimental.
|
||||
return {
|
||||
**checkpoint,
|
||||
"channel_values": {**checkpoint["channel_values"], **assembled},
|
||||
}
|
||||
|
||||
def materialize_checkpoint_tuple(
|
||||
self,
|
||||
value: CheckpointTuple,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> CheckpointTuple:
|
||||
"""Materialize incremental values for a single checkpoint tuple."""
|
||||
checkpoint = self.materialize_checkpoint(value.config, value.checkpoint, plan)
|
||||
if checkpoint is value.checkpoint:
|
||||
return value
|
||||
return value._replace(checkpoint=checkpoint)
|
||||
|
||||
def materialize_checkpoint_tuples(
|
||||
self,
|
||||
values: Sequence[CheckpointTuple],
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Sequence[CheckpointTuple]:
|
||||
"""Materialize incremental values for a batch of checkpoint tuples."""
|
||||
return [self.materialize_checkpoint_tuple(value, plan) for value in values]
|
||||
|
||||
async def amaterialize_checkpoint(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Async materialization hook for saver-managed incremental values."""
|
||||
if plan is None or not plan.channels:
|
||||
return checkpoint
|
||||
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
current_checkpoint_id = checkpoint.get("id")
|
||||
|
||||
targets = [
|
||||
(spec, checkpoint["channel_values"][spec.name])
|
||||
for spec in plan.channels
|
||||
if spec.kind == "delta"
|
||||
and isinstance(checkpoint["channel_values"].get(spec.name), DeltaValue)
|
||||
]
|
||||
if not targets:
|
||||
return checkpoint
|
||||
|
||||
# Walks for independent channels can run concurrently — each has its own chain.
|
||||
results = await asyncio.gather(
|
||||
*(
|
||||
self._amaterialize_delta_channel(
|
||||
thread_id=thread_id,
|
||||
checkpoint_ns=checkpoint_ns,
|
||||
current_checkpoint_id=current_checkpoint_id,
|
||||
channel=spec.name,
|
||||
value=value,
|
||||
)
|
||||
for spec, value in targets
|
||||
)
|
||||
)
|
||||
assembled = {
|
||||
spec.name: result
|
||||
for (spec, _), result in zip(targets, results, strict=True)
|
||||
if result is not None
|
||||
}
|
||||
|
||||
if not assembled:
|
||||
return checkpoint
|
||||
|
||||
return {
|
||||
**checkpoint,
|
||||
"channel_values": {**checkpoint["channel_values"], **assembled},
|
||||
}
|
||||
|
||||
async def amaterialize_checkpoint_tuple(
|
||||
self,
|
||||
value: CheckpointTuple,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> CheckpointTuple:
|
||||
"""Async materialization hook for a single checkpoint tuple."""
|
||||
checkpoint = await self.amaterialize_checkpoint(
|
||||
value.config, value.checkpoint, plan
|
||||
)
|
||||
if checkpoint is value.checkpoint:
|
||||
return value
|
||||
return value._replace(checkpoint=checkpoint)
|
||||
|
||||
async def amaterialize_checkpoint_tuples(
|
||||
self,
|
||||
values: Sequence[CheckpointTuple],
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Sequence[CheckpointTuple]:
|
||||
"""Async materialization hook for a batch of checkpoint tuples."""
|
||||
return [
|
||||
await self.amaterialize_checkpoint_tuple(value, plan) for value in values
|
||||
]
|
||||
|
||||
def get_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
"""Look up a single channel blob by checkpoint ID + channel name.
|
||||
|
||||
Returns NotImplemented if this saver does not support efficient
|
||||
per-channel-version blob lookup. The pregel layer will fall back to
|
||||
get_tuple() traversal in that case.
|
||||
|
||||
Savers with a dedicated blob store (InMemorySaver, PostgresSaver)
|
||||
should override this for O(1) performance.
|
||||
"""
|
||||
# Guard against re-entrant calls: when get_tuple() triggers
|
||||
# reconstruction which calls get_tuple() again, the inner call
|
||||
# short-circuits here.
|
||||
if _DELTA_RECONSTRUCTION.get():
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
return NotImplemented
|
||||
|
||||
token = _DELTA_RECONSTRUCTION.set(True)
|
||||
try:
|
||||
collected: list[PendingWrite] = [] # newest first; reversed at the end
|
||||
target_tuple = self.get_tuple(config)
|
||||
cursor_config: RunnableConfig | None = (
|
||||
target_tuple.parent_config if target_tuple else None
|
||||
)
|
||||
while cursor_config is not None:
|
||||
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)
|
||||
if tup.pending_writes:
|
||||
# Within a superstep, pending_writes are oldest→newest;
|
||||
# reverse to scan newest-first.
|
||||
for write in reversed(tup.pending_writes):
|
||||
if write[1] != channel:
|
||||
continue
|
||||
collected.append(write)
|
||||
cursor_config = tup.parent_config
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
|
||||
finally:
|
||||
_DELTA_RECONSTRUCTION.reset(token)
|
||||
async def aget_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
"""Look up a single channel blob by checkpoint ID + channel name (async).
|
||||
|
||||
async def _aget_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Async version of `_get_channel_writes_history`. See docstring there."""
|
||||
if _DELTA_RECONSTRUCTION.get():
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
Returns NotImplemented if this saver does not support efficient
|
||||
per-channel-version blob lookup. The pregel layer will fall back to
|
||||
aget_tuple() traversal in that case.
|
||||
|
||||
token = _DELTA_RECONSTRUCTION.set(True)
|
||||
try:
|
||||
collected: list[PendingWrite] = []
|
||||
target_tuple = await self.aget_tuple(config)
|
||||
cursor_config: RunnableConfig | None = (
|
||||
target_tuple.parent_config if target_tuple else None
|
||||
)
|
||||
while cursor_config is not None:
|
||||
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)
|
||||
if tup.pending_writes:
|
||||
for write in reversed(tup.pending_writes):
|
||||
if write[1] != channel:
|
||||
continue
|
||||
collected.append(write)
|
||||
cursor_config = tup.parent_config
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
|
||||
finally:
|
||||
_DELTA_RECONSTRUCTION.reset(token)
|
||||
Savers with a dedicated blob store (InMemorySaver, PostgresSaver)
|
||||
should override this for O(1) performance.
|
||||
"""
|
||||
return NotImplemented
|
||||
|
||||
def get_next_version(self, current: V | None, channel: None) -> V:
|
||||
"""Generate the next version ID for a channel.
|
||||
@@ -643,6 +696,138 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
clone.serde = maybe_add_typed_methods(serde)
|
||||
return clone
|
||||
|
||||
def _materialize_delta_channel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
current_checkpoint_id: str | None,
|
||||
channel: str,
|
||||
value: DeltaValue,
|
||||
) -> DeltaChainValue | None:
|
||||
chain_deltas: list[list[Any]] = []
|
||||
base: list[Any] | None = None
|
||||
cursor: DeltaValue = value
|
||||
visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set()
|
||||
|
||||
while True:
|
||||
chain_deltas.append(cursor.delta)
|
||||
prev_id = cursor.prev_checkpoint_id
|
||||
if prev_id is None:
|
||||
break
|
||||
if prev_id in visited:
|
||||
logger.warning(
|
||||
"DeltaChannel chain cycle at checkpoint %r for channel %r; breaking",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
visited.add(prev_id)
|
||||
|
||||
blob = self.get_channel_blob(thread_id, checkpoint_ns, prev_id, channel)
|
||||
if blob is not NotImplemented:
|
||||
if isinstance(blob, DeltaValue):
|
||||
cursor = blob
|
||||
continue
|
||||
base = blob
|
||||
break
|
||||
|
||||
parent_config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": prev_id,
|
||||
}
|
||||
}
|
||||
parent_tuple = self.get_tuple(parent_config)
|
||||
if parent_tuple is None:
|
||||
logger.warning(
|
||||
"DeltaChannel chain broken: checkpoint %r not found for channel %r",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
prev_val = parent_tuple.checkpoint["channel_values"].get(
|
||||
channel, _MISSING_SENTINEL
|
||||
)
|
||||
if prev_val is _MISSING_SENTINEL:
|
||||
break
|
||||
if isinstance(prev_val, DeltaValue):
|
||||
cursor = prev_val
|
||||
else:
|
||||
base = prev_val
|
||||
break
|
||||
|
||||
chain_deltas.reverse()
|
||||
return DeltaChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
async def _amaterialize_delta_channel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
current_checkpoint_id: str | None,
|
||||
channel: str,
|
||||
value: DeltaValue,
|
||||
) -> DeltaChainValue | None:
|
||||
chain_deltas: list[list[Any]] = []
|
||||
base: list[Any] | None = None
|
||||
cursor: DeltaValue = value
|
||||
visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set()
|
||||
|
||||
while True:
|
||||
chain_deltas.append(cursor.delta)
|
||||
prev_id = cursor.prev_checkpoint_id
|
||||
if prev_id is None:
|
||||
break
|
||||
if prev_id in visited:
|
||||
logger.warning(
|
||||
"DeltaChannel chain cycle at checkpoint %r for channel %r; breaking",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
visited.add(prev_id)
|
||||
|
||||
blob = await self.aget_channel_blob(
|
||||
thread_id, checkpoint_ns, prev_id, channel
|
||||
)
|
||||
if blob is not NotImplemented:
|
||||
if isinstance(blob, DeltaValue):
|
||||
cursor = blob
|
||||
continue
|
||||
base = blob
|
||||
break
|
||||
|
||||
parent_config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": prev_id,
|
||||
}
|
||||
}
|
||||
parent_tuple = await self.aget_tuple(parent_config)
|
||||
if parent_tuple is None:
|
||||
logger.warning(
|
||||
"DeltaChannel chain broken: checkpoint %r not found for channel %r",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
prev_val = parent_tuple.checkpoint["channel_values"].get(
|
||||
channel, _MISSING_SENTINEL
|
||||
)
|
||||
if prev_val is _MISSING_SENTINEL:
|
||||
break
|
||||
if isinstance(prev_val, DeltaValue):
|
||||
cursor = prev_val
|
||||
else:
|
||||
base = prev_val
|
||||
break
|
||||
|
||||
chain_deltas.reverse()
|
||||
return DeltaChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
|
||||
def _with_msgpack_allowlist(
|
||||
serde: SerializerProtocol, extra_allowlist: Collection[tuple[str, ...]]
|
||||
|
||||
@@ -9,25 +9,21 @@ from collections import defaultdict
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
|
||||
from types import TracebackType
|
||||
from typing import Any, cast
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
PendingWrite,
|
||||
SerializerProtocol,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -125,114 +121,50 @@ class InMemorySaver(
|
||||
return self.stack.__exit__(__exc_type, __exc_value, __traceback)
|
||||
|
||||
def _load_blobs(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
versions: ChannelVersions,
|
||||
self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions
|
||||
) -> dict[str, Any]:
|
||||
result: dict[str, Any] = {}
|
||||
for k, ver in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, ver)
|
||||
channel_values: dict[str, Any] = {}
|
||||
for k, v in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, v)
|
||||
if kk not in self.blobs:
|
||||
continue
|
||||
vv = self.blobs[kk]
|
||||
if vv[0] == "empty":
|
||||
continue
|
||||
result[k] = self.serde.loads_typed(vv)
|
||||
return result
|
||||
if vv[0] != "empty":
|
||||
channel_values[k] = self.serde.loads_typed(vv)
|
||||
return channel_values
|
||||
|
||||
def _get_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = config["configurable"].get("checkpoint_id", "")
|
||||
def get_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
"""Fast-path blob lookup: checkpoint → channel version → blob."""
|
||||
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
|
||||
# Walk the parent chain newest→oldest. Skip the target itself —
|
||||
# writes stored AT `checkpoint_id` are pending for the next step
|
||||
# (pregel applies them via `apply_writes`; they aren't part of the
|
||||
# snapshot value AT `checkpoint_id`).
|
||||
chain: list[str] = []
|
||||
target_entry = ns_storage.get(checkpoint_id)
|
||||
current: str | None = target_entry[2] if target_entry is not None else None
|
||||
while current is not None:
|
||||
entry = ns_storage.get(current)
|
||||
if entry is None:
|
||||
break
|
||||
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).
|
||||
collected: list[PendingWrite] = [] # newest first
|
||||
for cp_id in chain: # newest → oldest
|
||||
entry = ns_storage.get(cp_id)
|
||||
if entry is not None:
|
||||
ckpt = self.serde.loads_typed(entry[0])
|
||||
ver = ckpt.get("channel_versions", {}).get(channel)
|
||||
if ver is not None:
|
||||
blob_entry = self.blobs.get(
|
||||
(thread_id, checkpoint_ns, channel, ver)
|
||||
)
|
||||
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.
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
)
|
||||
entry = ns_storage.get(checkpoint_id)
|
||||
if entry is None:
|
||||
return NotImplemented
|
||||
checkpoint = self.serde.loads_typed(entry[0])
|
||||
version = checkpoint["channel_versions"].get(channel)
|
||||
if version is None:
|
||||
return NotImplemented
|
||||
kk = (thread_id, checkpoint_ns, channel, version)
|
||||
if kk not in self.blobs:
|
||||
return NotImplemented
|
||||
vv = self.blobs[kk]
|
||||
if vv[0] == "empty":
|
||||
return NotImplemented
|
||||
return self.serde.loads_typed(vv)
|
||||
|
||||
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)
|
||||
|
||||
async def _aget_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
return self._get_channel_writes_history(config, channel)
|
||||
async def aget_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
return self.get_channel_blob(thread_id, checkpoint_ns, checkpoint_id, channel)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the in-memory storage.
|
||||
@@ -255,16 +187,13 @@ class InMemorySaver(
|
||||
checkpoint, metadata, parent_checkpoint_id = saved
|
||||
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
|
||||
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
|
||||
channel_values = self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
)
|
||||
return CheckpointTuple(
|
||||
config=config,
|
||||
checkpoint={
|
||||
**checkpoint_,
|
||||
"channel_values": channel_values,
|
||||
"channel_values": self._load_blobs(
|
||||
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
|
||||
),
|
||||
},
|
||||
metadata=self.serde.loads_typed(metadata),
|
||||
pending_writes=[
|
||||
@@ -288,26 +217,19 @@ class InMemorySaver(
|
||||
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
|
||||
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
|
||||
checkpoint_ = self.serde.loads_typed(checkpoint)
|
||||
resolved_config = cast(
|
||||
RunnableConfig,
|
||||
{
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
)
|
||||
channel_values = self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
)
|
||||
return CheckpointTuple(
|
||||
config=resolved_config,
|
||||
checkpoint={
|
||||
**checkpoint_,
|
||||
"channel_values": channel_values,
|
||||
"channel_values": self._load_blobs(
|
||||
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
|
||||
),
|
||||
},
|
||||
metadata=self.serde.loads_typed(metadata),
|
||||
pending_writes=[
|
||||
@@ -402,27 +324,21 @@ class InMemorySaver(
|
||||
|
||||
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
|
||||
|
||||
list_config = cast(
|
||||
RunnableConfig,
|
||||
{
|
||||
yield CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
)
|
||||
channel_values = self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
)
|
||||
|
||||
yield CheckpointTuple(
|
||||
config=list_config,
|
||||
checkpoint={
|
||||
**checkpoint_,
|
||||
"channel_values": channel_values,
|
||||
"channel_values": self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
),
|
||||
},
|
||||
metadata=metadata,
|
||||
parent_config=(
|
||||
@@ -525,101 +441,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.
|
||||
|
||||
|
||||
@@ -80,6 +80,8 @@ SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
|
||||
("langgraph.types", "Overwrite"),
|
||||
("langgraph.store.base", "Item"),
|
||||
("langgraph.store.base", "GetOp"),
|
||||
# DeltaChannel checkpoint value type
|
||||
("langgraph.checkpoint.base", "DeltaValue"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -33,39 +33,24 @@ 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 SendProtocol
|
||||
from langgraph.store.base import Item
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.checkpoint.serde._msgpack import (
|
||||
AllowedMsgpackModules,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import SendProtocol
|
||||
|
||||
LC_REVIVER = Reviver()
|
||||
EMPTY_BYTES = b""
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Dedup log warnings across process lifetime; cap bounds state if types are
|
||||
# dynamically generated (also acts as a circuit breaker on warning volume).
|
||||
# Dedup is best-effort: racing threads may each emit once for the same key,
|
||||
# and warnings are silently dropped once _MAX_WARNED_TYPES is reached.
|
||||
_MAX_WARNED_TYPES = 1000
|
||||
_warned_unregistered_types: set[tuple[str, str]] = set()
|
||||
_warned_blocked_types: set[tuple[str, str]] = set()
|
||||
|
||||
def _is_delta_value(obj: Any) -> bool:
|
||||
from langgraph.checkpoint.base import DeltaValue # lazy import avoids circular dep
|
||||
|
||||
def _warn_once(
|
||||
seen: set[tuple[str, str]], key: tuple[str, str], msg: str, *args: object
|
||||
) -> None:
|
||||
if key in seen or len(seen) >= _MAX_WARNED_TYPES:
|
||||
return
|
||||
seen.add(key)
|
||||
logger.warning(msg, *args)
|
||||
return isinstance(obj, DeltaValue)
|
||||
|
||||
|
||||
class JsonPlusSerializer(SerializerProtocol):
|
||||
@@ -260,6 +245,8 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return "bytes", obj
|
||||
elif isinstance(obj, bytearray):
|
||||
return "bytearray", obj
|
||||
elif _is_delta_value(obj):
|
||||
return "delta", _msgpack_enc({"d": obj.delta, "c": obj.prev_checkpoint_id})
|
||||
else:
|
||||
try:
|
||||
return "msgpack", _msgpack_enc(obj)
|
||||
@@ -282,6 +269,13 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
elif type_ == "delta":
|
||||
from langgraph.checkpoint.base import DeltaValue # lazy import
|
||||
|
||||
raw = ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
return DeltaValue(delta=raw["d"], prev_checkpoint_id=raw.get("c"))
|
||||
elif self.pickle_fallback and type_ == "pickle":
|
||||
return pickle.loads(data_)
|
||||
else:
|
||||
@@ -297,16 +291,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(
|
||||
@@ -561,9 +549,7 @@ def _create_msgpack_ext_hook(
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
_warn_once(
|
||||
_warned_unregistered_types,
|
||||
key,
|
||||
logger.warning(
|
||||
"Deserializing unregistered type %s.%s from checkpoint. "
|
||||
"This will be blocked in a future version. "
|
||||
"Set LANGGRAPH_STRICT_MSGPACK=true to block now, or add "
|
||||
@@ -585,9 +571,7 @@ def _create_msgpack_ext_hook(
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
_warn_once(
|
||||
_warned_blocked_types,
|
||||
key,
|
||||
logger.warning(
|
||||
"Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
|
||||
"Add to allowed_msgpack_modules to allow: [(%r, %r)]",
|
||||
module,
|
||||
@@ -620,15 +604,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,
|
||||
@@ -15,39 +14,6 @@ INTERRUPT = "__interrupt__"
|
||||
RESUME = "__resume__"
|
||||
TASKS = "__pregel_tasks"
|
||||
|
||||
|
||||
class _DeltaSentinel:
|
||||
"""Singleton marker stored (as zero bytes) in checkpoint_blobs for a
|
||||
DeltaChannel field. The actual per-step writes live in checkpoint_writes
|
||||
and are replayed through the reducer at load time.
|
||||
|
||||
Compare with `is DELTA_SENTINEL` — `loads_typed` always returns the same
|
||||
module-level instance.
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "DELTA_SENTINEL"
|
||||
|
||||
|
||||
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")
|
||||
|
||||
@@ -29,8 +29,6 @@ from langgraph.checkpoint.serde.jsonplus import (
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
|
||||
|
||||
@@ -104,13 +102,6 @@ def test_msgpack_method_pathlib_blocked_encrypted_strict(
|
||||
class TestEncryptedSerializerMsgpackAllowlist:
|
||||
"""Test msgpack allowlist behavior through EncryptedSerializer."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_warned_types(self) -> None:
|
||||
# Warning dedup state is process-global; reset per-test so each case
|
||||
# sees a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
|
||||
def test_safe_types_no_warning(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Test safe types deserialize without warnings through encryption."""
|
||||
serde = _make_encrypted_serde()
|
||||
|
||||
@@ -35,8 +35,6 @@ from langgraph.checkpoint.serde.jsonplus import (
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_msgpack_ext_hook_to_json,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
from langgraph.store.base import Item
|
||||
|
||||
@@ -582,14 +580,6 @@ def test_msgpack_safe_types_no_warning(caplog: pytest.LogCaptureFixture) -> None
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_warned_types() -> None:
|
||||
# Warning dedup state is process-global; reset per-test so each case sees
|
||||
# a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
|
||||
|
||||
def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Pydantic models not in allowlist should log warning but still deserialize."""
|
||||
current = _lg_msgpack.STRICT_MSGPACK_ENABLED
|
||||
@@ -605,12 +595,6 @@ def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) ->
|
||||
assert "unregistered type" in caplog.text.lower()
|
||||
assert "allowed_msgpack_modules" in caplog.text
|
||||
assert result == obj
|
||||
|
||||
# Second deserialization of the same type should NOT produce another warning
|
||||
caplog.clear()
|
||||
result2 = serde.loads_typed(dumped)
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert result2 == obj
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = current
|
||||
|
||||
|
||||
@@ -655,6 +639,7 @@ def test_msgpack_allowlist_silences_warning(caplog: pytest.LogCaptureFixture) ->
|
||||
|
||||
def test_msgpack_none_blocks_unregistered(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""allowed_msgpack_modules=None should block unregistered types."""
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
@@ -672,6 +657,7 @@ def test_msgpack_allowlist_blocks_non_listed(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Allowlists should block unregistered types even if msgpack is enabled."""
|
||||
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")]
|
||||
)
|
||||
@@ -999,14 +985,29 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_delta_sentinel_serde_round_trip() -> None:
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
def test_delta_value_serde_round_trip() -> None:
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
type_tag, blob = serde.dumps_typed(DELTA_SENTINEL)
|
||||
# Zero-byte "delta" tag — no allowlist change needed.
|
||||
original = DeltaValue(
|
||||
delta=[{"type": "human", "content": "hi"}], prev_checkpoint_id="abc-123"
|
||||
)
|
||||
type_tag, blob = serde.dumps_typed(original)
|
||||
assert type_tag == "delta"
|
||||
assert blob == b""
|
||||
loaded = serde.loads_typed((type_tag, blob))
|
||||
assert loaded is DELTA_SENTINEL
|
||||
assert isinstance(loaded, DeltaValue)
|
||||
assert loaded.delta == original.delta
|
||||
assert loaded.prev_checkpoint_id == "abc-123"
|
||||
|
||||
|
||||
def test_delta_value_serde_chain_root() -> None:
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
original = DeltaValue(delta=[], prev_checkpoint_id=None)
|
||||
type_tag, blob = serde.dumps_typed(original)
|
||||
loaded = serde.loads_typed((type_tag, blob))
|
||||
assert isinstance(loaded, DeltaValue)
|
||||
assert loaded.prev_checkpoint_id is None
|
||||
|
||||
@@ -6,32 +6,19 @@ from langchain_core.runnables import RunnableConfig
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.jsonplus import (
|
||||
JsonPlusSerializer,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
|
||||
class MemoryPydantic(BaseModel):
|
||||
foo: str
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_warned_types() -> None:
|
||||
# Warning dedup state is process-global; reset per-test so each case sees
|
||||
# a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
|
||||
|
||||
class TestMemorySaver:
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self) -> None:
|
||||
@@ -209,6 +196,8 @@ class TestMemorySaver:
|
||||
|
||||
|
||||
async def test_memory_saver() -> None:
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
memory_saver = InMemorySaver()
|
||||
assert isinstance(memory_saver, InMemorySaver)
|
||||
|
||||
@@ -322,489 +311,33 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
|
||||
|
||||
|
||||
class TestInMemorySaverDeltaChannel:
|
||||
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
|
||||
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
|
||||
def test_get_channel_blob(self) -> None:
|
||||
"""get_channel_blob returns the deserialized blob for a checkpoint+channel."""
|
||||
from langgraph.checkpoint.base import DeltaValue, empty_checkpoint
|
||||
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
v1 = "00000000000000000000000000000001.0000000000000000"
|
||||
version = "00000000000000000000000000000001.0000000000000000"
|
||||
delta = DeltaValue(delta=[{"content": "hi"}], prev_checkpoint_id=None)
|
||||
saver.blobs[(thread_id, ns, channel, version)] = serde.dumps_typed(delta)
|
||||
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL)
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp1["channel_versions"][channel] = v1
|
||||
cp = empty_checkpoint()
|
||||
cp["id"] = "cp1"
|
||||
cp["channel_versions"][channel] = version
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
"cp1": (serde.dumps_typed(cp), serde.dumps_typed({}), None)
|
||||
}
|
||||
|
||||
result = saver._load_blobs(thread_id, ns, {channel: v1})
|
||||
assert channel in result
|
||||
assert result[channel] is DELTA_SENTINEL
|
||||
result = saver.get_channel_blob(thread_id, ns, "cp1", channel)
|
||||
assert isinstance(result, DeltaValue)
|
||||
assert result.delta == [{"content": "hi"}]
|
||||
assert result.prev_checkpoint_id is None
|
||||
|
||||
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
|
||||
"""_get_channel_writes_history collects ancestor writes oldest→newest,
|
||||
and excludes writes stored at the target checkpoint itself (those are
|
||||
pending writes for the next step, applied separately by pregel)."""
|
||||
def test_get_channel_blob_missing(self) -> None:
|
||||
"""get_channel_blob returns NotImplemented when checkpoint or channel not found."""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "cp2"
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
||||
}
|
||||
# Writes stored at cp1 produced the cp1 snapshot; part of history.
|
||||
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "hi"}),
|
||||
"",
|
||||
)
|
||||
# Writes stored at cp2 are pending — they will produce cp3 when the
|
||||
# step that loaded cp2 completes. They MUST NOT appear in the
|
||||
# reconstructed snapshot value at cp2.
|
||||
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "pending"}),
|
||||
"",
|
||||
)
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": "cp2",
|
||||
}
|
||||
}
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == [{"content": "hi"}]
|
||||
|
||||
def test_get_channel_writes_at_root_returns_empty(self) -> None:
|
||||
"""Reconstructing the root checkpoint's state: no ancestors → []."""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
}
|
||||
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "pending"}),
|
||||
"",
|
||||
)
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": "cp1",
|
||||
}
|
||||
}
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
assert result.writes == []
|
||||
|
||||
|
||||
class TestBaseFallbackGetChannelWrites:
|
||||
"""Exercises the `BaseCheckpointSaver._get_channel_writes_history` default
|
||||
implementation — the path third-party savers inherit when they don't
|
||||
override `_get_channel_writes_history` themselves.
|
||||
|
||||
Regression guard for a bug where the fallback passed the caller's config
|
||||
(with `checkpoint_id`) straight to `self.list()`, which most savers
|
||||
collapse to a single row — causing the fallback to return `[]`.
|
||||
"""
|
||||
|
||||
def _build_saver_with_chain(self) -> tuple[InMemorySaver, str, str]:
|
||||
"""Build an InMemorySaver with a 3-checkpoint chain and per-step writes
|
||||
for a `messages` channel.
|
||||
|
||||
Returns `(saver, thread_id, namespace)`. The saver subclass deletes the
|
||||
InMemorySaver override so the base class fallback is exercised.
|
||||
"""
|
||||
|
||||
class _ThirdPartyStyleSaver(InMemorySaver):
|
||||
_get_channel_writes_history = (
|
||||
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
_aget_channel_writes_history = (
|
||||
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
saver = _ThirdPartyStyleSaver()
|
||||
serde = JsonPlusSerializer()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
cp0 = empty_checkpoint()
|
||||
cp0["id"] = "00000000000000000000000000000001.0000000000000000"
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "00000000000000000000000000000002.0000000000000000"
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "00000000000000000000000000000003.0000000000000000"
|
||||
saver.storage[thread_id][ns] = {
|
||||
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
|
||||
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
|
||||
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]),
|
||||
}
|
||||
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
|
||||
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "first"}),
|
||||
"",
|
||||
)
|
||||
saver.writes[(thread_id, ns, cp1["id"])][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "second"}),
|
||||
"",
|
||||
)
|
||||
return saver, thread_id, ns
|
||||
|
||||
def test_fallback_returns_ancestor_writes_oldest_first(self) -> None:
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._get_channel_writes_history(config, "messages")
|
||||
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == [{"content": "first"}, {"content": "second"}]
|
||||
|
||||
async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None:
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
result = await saver._aget_channel_writes_history(config, "messages")
|
||||
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == [{"content": "first"}, {"content": "second"}]
|
||||
|
||||
async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None:
|
||||
"""Regression: the re-entrancy guard must be task-local, not thread-local.
|
||||
|
||||
Two concurrent `_aget_channel_writes_history` calls on the same
|
||||
event-loop thread must each see their full reconstructed writes. A
|
||||
`threading.local()` guard would let whichever task set it first
|
||||
short-circuit the other to `writes=[]`.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
|
||||
# Force the two tasks to interleave across the `set(True)` boundary:
|
||||
# each `aget_tuple` yields control, so if the guard were thread-local
|
||||
# the second task would observe `active=True` set by the first.
|
||||
orig_aget_tuple = saver.aget_tuple
|
||||
|
||||
async def slow_aget_tuple(config: RunnableConfig) -> Any:
|
||||
await asyncio.sleep(0)
|
||||
return await orig_aget_tuple(config)
|
||||
|
||||
saver.aget_tuple = slow_aget_tuple # type: ignore[method-assign]
|
||||
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
results = await asyncio.gather(
|
||||
saver._aget_channel_writes_history(config, "messages"),
|
||||
saver._aget_channel_writes_history(config, "messages"),
|
||||
)
|
||||
|
||||
expected_values = [{"content": "first"}, {"content": "second"}]
|
||||
for result in results:
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == expected_values
|
||||
|
||||
|
||||
class TestPreDeltaBlobTerminator:
|
||||
"""Verify the pre-delta blob terminator: when the ancestor walk hits a
|
||||
checkpoint whose blob for the channel is a real value (not
|
||||
DELTA_SENTINEL), reconstruction seeds from it and stops. This guards
|
||||
|
||||
* back-compat: a thread written by pre-delta code, then extended under
|
||||
delta — reconstruction must return the correct value without walking
|
||||
past the last pre-delta ancestor;
|
||||
* perf: without the terminator, every reconstruct-after-migration would
|
||||
walk all the way to the thread root.
|
||||
"""
|
||||
|
||||
def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]:
|
||||
"""Three-checkpoint chain: cp1 (pre-delta, blob=[A]), cp2 (delta,
|
||||
write=B), cp3 (delta, write=C). Reconstructing at cp3 must yield
|
||||
seed=[A] + writes=[B, C].
|
||||
|
||||
Returns `(saver, thread_id, ns, channel, cp3_id)`.
|
||||
"""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
v1 = "00000000000000000000000000000001.0"
|
||||
v2 = "00000000000000000000000000000002.0"
|
||||
v3 = "00000000000000000000000000000003.0"
|
||||
|
||||
# Pre-delta: cp1 stored a real blob for the channel.
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"])
|
||||
# Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes.
|
||||
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL)
|
||||
saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL)
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp1["channel_versions"][channel] = v1
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "cp2"
|
||||
cp2["channel_versions"][channel] = v2
|
||||
cp3 = empty_checkpoint()
|
||||
cp3["id"] = "cp3"
|
||||
cp3["channel_versions"][channel] = v3
|
||||
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
||||
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"),
|
||||
}
|
||||
# Write under cp1 would be from the pre-delta era and MUST be ignored
|
||||
# (the blob already captures it). We add one and assert it is not
|
||||
# folded into the reconstructed result.
|
||||
saver.writes[(thread_id, ns, "cp1")][("task0", 0)] = (
|
||||
"task0",
|
||||
channel,
|
||||
serde.dumps_typed("PRE-DELTA-WRITE"),
|
||||
"",
|
||||
)
|
||||
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed("B"),
|
||||
"",
|
||||
)
|
||||
saver.writes[(thread_id, ns, "cp3")][("task3", 0)] = (
|
||||
"task3",
|
||||
channel,
|
||||
serde.dumps_typed("PENDING-AT-TARGET"),
|
||||
"",
|
||||
)
|
||||
return saver, thread_id, ns, channel, "cp3"
|
||||
|
||||
def test_seed_from_pre_delta_ancestor_blob(self) -> None:
|
||||
saver, thread_id, ns, channel, target = self._build_mixed_thread()
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
|
||||
# Seed came from the pre-delta blob at cp1.
|
||||
assert result.seed == ["A"]
|
||||
# Delta-era writes from cp2 replay through the reducer on top of seed.
|
||||
# cp3 is the target — its own write is pending for the NEXT step and
|
||||
# must be excluded.
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == ["B"]
|
||||
|
||||
def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None:
|
||||
"""Writes stored at the pre-delta ancestor itself must not be replayed
|
||||
(the blob subsumes them)."""
|
||||
saver, thread_id, ns, channel, target = self._build_mixed_thread()
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
|
||||
values = [v for _, _, v in result.writes]
|
||||
# The pre-delta write under cp1 must not appear (the blob subsumes it).
|
||||
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.get_channel_blob("t1", "", "no-such-cp", "messages") is NotImplemented
|
||||
)
|
||||
|
||||
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]}
|
||||
|
||||
@@ -5,5 +5,5 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.1.14"
|
||||
"langchain-openai==1.0.1"
|
||||
]
|
||||
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.1.14",
|
||||
"langchain-openai==1.0.0a2",
|
||||
"langchain-anthropic==1.0.0a5",
|
||||
"langgraph==1.1.5"
|
||||
]
|
||||
|
||||
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.1.14",
|
||||
"langchain-openai==1.0.0a2",
|
||||
"langgraph==1.1.2",
|
||||
"langchain_community>=0.3.0",
|
||||
]
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.4.23"
|
||||
__version__ = "0.4.22"
|
||||
|
||||
@@ -23,7 +23,7 @@ dependencies = [
|
||||
path = "langgraph_cli/__init__.py"
|
||||
[project.optional-dependencies]
|
||||
inmem = [
|
||||
"langgraph-api>=0.5.35,<0.9.0 ; python_version >= '3.11'",
|
||||
"langgraph-api>=0.5.35,<0.8.0 ; python_version >= '3.11'",
|
||||
"langgraph-runtime-inmem>=0.7 ; python_version >= '3.11'",
|
||||
]
|
||||
|
||||
|
||||
Generated
+383
-465
File diff suppressed because it is too large
Load Diff
@@ -245,6 +245,15 @@ class _GraphCallbackManager(BaseCallbackManager):
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
def add_handler(
|
||||
self,
|
||||
handler: BaseCallbackHandler,
|
||||
inherit: bool = True, # noqa: FBT001,FBT002
|
||||
) -> None:
|
||||
if not isinstance(handler, GraphCallbackHandler):
|
||||
raise TypeError("handlers must inherit GraphCallbackHandler")
|
||||
super().add_handler(handler, inherit=inherit)
|
||||
|
||||
def copy(
|
||||
self,
|
||||
*,
|
||||
@@ -312,6 +321,15 @@ class _AsyncGraphCallbackManager(BaseCallbackManager):
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
def add_handler(
|
||||
self,
|
||||
handler: BaseCallbackHandler,
|
||||
inherit: bool = True, # noqa: FBT001,FBT002
|
||||
) -> None:
|
||||
if not isinstance(handler, GraphCallbackHandler):
|
||||
raise TypeError("handlers must inherit GraphCallbackHandler")
|
||||
super().add_handler(handler, inherit=inherit)
|
||||
|
||||
def copy(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Generic, TypeVar
|
||||
from typing import Any, Generic, Literal, TypeVar
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
@@ -119,3 +119,17 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
|
||||
Returns `True` if the channel was updated, `False` otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
|
||||
"""Called after checkpoint() with the assigned version, and after
|
||||
from_checkpoint() with the current channel version.
|
||||
|
||||
No-op by default. Override in channels that track their own version
|
||||
for incremental checkpointing (e.g. DeltaChannel).
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
def checkpoint_hydration_kind(self) -> Literal["delta"] | None:
|
||||
"""Return the saver hydration kind for this channel, if any."""
|
||||
return None
|
||||
|
||||
@@ -1,78 +1,97 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy as _copy
|
||||
import collections.abc
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
from copy import copy
|
||||
from typing import Any, Generic, Literal
|
||||
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
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
|
||||
from langgraph.channels.binop import _get_overwrite, _strip_extras
|
||||
from langgraph.errors import EmptyChannelError
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
def _empty(typ: Any) -> Any:
|
||||
def _copy_value(value: Any) -> Any:
|
||||
if value is MISSING:
|
||||
return value
|
||||
try:
|
||||
return typ()
|
||||
except Exception:
|
||||
return []
|
||||
return value.copy()
|
||||
except AttributeError:
|
||||
return copy(value)
|
||||
|
||||
|
||||
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[list[Value], Value, DeltaValue]):
|
||||
"""A channel that stores only per-step write deltas in checkpoints.
|
||||
|
||||
Reconstructs the full accumulated list at load time by replaying the
|
||||
chain of deltas through the operator. Use with append-style reducers
|
||||
(e.g. `add_messages`) on long-running threads to reduce checkpoint
|
||||
storage from O(N²) to O(N).
|
||||
|
||||
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
"""Fold-reducer channel with configurable snapshot cadence.
|
||||
Works with all checkpointers. Savers with a dedicated blob store
|
||||
(InMemorySaver, PostgresSaver) use an O(1) fast-path per chain step;
|
||||
all others (SQLite, MongoDB, etc.) fall back to get_tuple traversal.
|
||||
|
||||
`snapshot_frequency=None` (default): pure delta — stores only
|
||||
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes.
|
||||
Use `snapshot_every=N` to cap chain traversal depth at N steps. Every N
|
||||
steps a full snapshot is written as the chain root; subsequent deltas
|
||||
chain back to it, so `get_state` / reload never traverses more than N
|
||||
checkpoints regardless of thread length. Recommended for savers without
|
||||
a dedicated blob store.
|
||||
|
||||
`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.
|
||||
Usage::
|
||||
|
||||
Parameters:
|
||||
operator: Binary reducer `(Value, Value) -> Value`.
|
||||
snapshot_frequency: Every Nth pregel step writes a snapshot blob.
|
||||
`None` (default) = pure delta, never snapshot.
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
|
||||
# Cap reconstruction depth (recommended for SQLite / MongoDB savers):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages, snapshot_every=50)]
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "operator", "snapshot_frequency")
|
||||
value: Value | Any
|
||||
__slots__ = (
|
||||
"value",
|
||||
"operator",
|
||||
"snapshot_every",
|
||||
"_pending",
|
||||
"_base_version",
|
||||
"_last_checkpoint_id",
|
||||
"_overwritten",
|
||||
"_steps_since_snapshot",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[Any, Any], Any],
|
||||
operator: Callable[[list[Value], Any], list[Value]],
|
||||
typ: type = list,
|
||||
*,
|
||||
snapshot_frequency: int | None = None,
|
||||
snapshot_every: int | None = None,
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
typ = _strip_extras(typ)
|
||||
if typ in (
|
||||
collections.abc.Sequence,
|
||||
collections.abc.MutableSequence,
|
||||
):
|
||||
typ = list
|
||||
super().__init__(typ)
|
||||
self.operator = operator
|
||||
self.snapshot_frequency = snapshot_frequency
|
||||
self.value: Any = []
|
||||
self.snapshot_every = snapshot_every
|
||||
try:
|
||||
self.value: list[Value] = typ()
|
||||
except Exception:
|
||||
self.value = []
|
||||
self._pending: list[Any] = []
|
||||
self._base_version: str | None = None
|
||||
self._last_checkpoint_id: str | None = None
|
||||
self._overwritten: bool = False
|
||||
self._steps_since_snapshot: int = 0
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
return False
|
||||
if self.snapshot_frequency != other.snapshot_frequency:
|
||||
if self.snapshot_every != other.snapshot_every:
|
||||
return False
|
||||
if (
|
||||
self.operator.__name__ != "<lambda>"
|
||||
@@ -83,85 +102,96 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
|
||||
@property
|
||||
def ValueType(self) -> Any:
|
||||
return self.typ
|
||||
return list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ
|
||||
return self.typ | list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
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
|
||||
@property
|
||||
def checkpoint_hydration_kind(self) -> Literal["delta"]:
|
||||
return "delta"
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = self._clone_empty()
|
||||
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
||||
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
|
||||
new.key = self.key
|
||||
new.value = _copy_value(self.value)
|
||||
new._pending = self._pending[:]
|
||||
new._base_version = self._base_version
|
||||
new._last_checkpoint_id = self._last_checkpoint_id
|
||||
new._overwritten = self._overwritten
|
||||
new._steps_since_snapshot = self._steps_since_snapshot
|
||||
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
|
||||
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
pass
|
||||
elif isinstance(checkpoint, DeltaChainValue):
|
||||
accumulated: list[Value] = (
|
||||
checkpoint.base if checkpoint.base is not None else new.typ()
|
||||
)
|
||||
for step_writes in checkpoint.deltas:
|
||||
for write in step_writes:
|
||||
accumulated = new.operator(accumulated, write)
|
||||
new.value = accumulated
|
||||
# Seed the counter from actual chain depth so rehydration fires at
|
||||
# the right time regardless of how many prior invocations there were.
|
||||
new._steps_since_snapshot = len(checkpoint.deltas)
|
||||
elif isinstance(checkpoint, DeltaValue):
|
||||
# Should never reach here — checkpoint hydration should assemble
|
||||
# DeltaValues into DeltaChainValue before calling from_checkpoint.
|
||||
raise AssertionError(
|
||||
"DeltaChannel.from_checkpoint received a raw DeltaValue. "
|
||||
"This is a bug in checkpoint hydration — chain assembly should "
|
||||
"have occurred before from_checkpoint was called."
|
||||
)
|
||||
else:
|
||||
new.value = checkpoint
|
||||
# Backwards compat: plain value from old BinaryOperatorAggregate checkpoint
|
||||
# or a full snapshot emitted by DeltaChannel.
|
||||
new.value = _copy_value(checkpoint)
|
||||
new._pending = []
|
||||
new._base_version = None # set by the subsequent after_checkpoint() call
|
||||
new._overwritten = False
|
||||
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)
|
||||
is_overwrite, overwrite_value = _get_overwrite(value)
|
||||
if is_overwrite:
|
||||
if seen_overwrite:
|
||||
from langgraph.errors import (
|
||||
ErrorCode,
|
||||
InvalidUpdateError,
|
||||
create_error_message,
|
||||
)
|
||||
|
||||
msg = create_error_message(
|
||||
message="Can receive only one Overwrite value per super-step.",
|
||||
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
|
||||
)
|
||||
raise InvalidUpdateError(msg)
|
||||
self.value = (
|
||||
_copy_value(overwrite_value)
|
||||
if overwrite_value is not None
|
||||
else self.typ()
|
||||
)
|
||||
self._pending = (
|
||||
[] if overwrite_value is None else [_copy_value(self.value)]
|
||||
)
|
||||
self._overwritten = True
|
||||
seen_overwrite = True
|
||||
elif seen_overwrite:
|
||||
continue
|
||||
self.value = self._apply_write(self.value, value)
|
||||
elif not seen_overwrite:
|
||||
base = self.typ() if self.value is MISSING else self.value
|
||||
self.value = self.operator(base, value)
|
||||
self._pending.append(value)
|
||||
return True
|
||||
|
||||
def get(self) -> Any:
|
||||
def get(self) -> list[Value]:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
@@ -170,12 +200,29 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
"""Return stored representation: always `DELTA_SENTINEL`.
|
||||
if (
|
||||
self.snapshot_every is not None
|
||||
and self._steps_since_snapshot >= self.snapshot_every
|
||||
):
|
||||
# Emit a full snapshot to cap chain depth at snapshot_every.
|
||||
# The saver stores this as a plain (non-diff) blob, so future
|
||||
# deltas will chain back to it and traversal depth resets to 1.
|
||||
return _copy_value(self.value)
|
||||
return DeltaValue(
|
||||
delta=self._pending[:],
|
||||
prev_checkpoint_id=None if self._overwritten else self._last_checkpoint_id,
|
||||
)
|
||||
|
||||
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
|
||||
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
|
||||
if version != self._base_version:
|
||||
if self._base_version is None:
|
||||
pass # First call after from_checkpoint — anchor without counting a step.
|
||||
elif self.snapshot_every is not None:
|
||||
if self._steps_since_snapshot >= self.snapshot_every:
|
||||
self._steps_since_snapshot = 0
|
||||
else:
|
||||
self._steps_since_snapshot += 1
|
||||
self._base_version = version
|
||||
self._last_checkpoint_id = checkpoint_id
|
||||
self._pending = []
|
||||
self._overwritten = False
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import collections.abc
|
||||
import inspect
|
||||
import logging
|
||||
import typing
|
||||
@@ -48,8 +47,7 @@ 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.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate, _strip_extras
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
|
||||
from langgraph.channels.named_barrier_value import (
|
||||
@@ -1084,7 +1082,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
CompiledStateGraph: The compiled `StateGraph`.
|
||||
"""
|
||||
checkpointer = ensure_valid_checkpointer(checkpointer)
|
||||
|
||||
serde_allowlist: set[tuple[str, ...]] | None = None
|
||||
if _serde.STRICT_MSGPACK_ENABLED:
|
||||
schema_types: list[type[Any]] = [
|
||||
@@ -1670,36 +1667,6 @@ 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__"):
|
||||
origin = typ.__origin__
|
||||
# Unwrap parameterized Required[X]/NotRequired[X] to X
|
||||
# (e.g. Annotated[NotRequired[dict[...]], ...]).
|
||||
if hasattr(origin, "__origin__") and origin.__origin__ in (
|
||||
Required,
|
||||
NotRequired,
|
||||
):
|
||||
origin = origin.__args__[0]
|
||||
outer = _strip_extras(origin)
|
||||
if outer in (
|
||||
collections.abc.Sequence,
|
||||
collections.abc.MutableSequence,
|
||||
):
|
||||
outer = list
|
||||
elif outer in (
|
||||
collections.abc.Mapping,
|
||||
collections.abc.MutableMapping,
|
||||
):
|
||||
outer = dict
|
||||
elif outer in (
|
||||
collections.abc.Set,
|
||||
collections.abc.MutableSet,
|
||||
):
|
||||
outer = set
|
||||
item.typ = outer
|
||||
try:
|
||||
item.value = outer()
|
||||
except Exception:
|
||||
item.value = []
|
||||
return item
|
||||
elif isclass(item) and issubclass(item, BaseChannel):
|
||||
# ex, Annotated[int, EphemeralValue, SomeOtherAnnotation]
|
||||
|
||||
@@ -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,22 +1,33 @@
|
||||
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 import (
|
||||
Checkpoint,
|
||||
CheckpointHydrationPlan,
|
||||
IncrementalChannelSpec,
|
||||
)
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
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 checkpoint_hydration_plan(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
) -> CheckpointHydrationPlan | None:
|
||||
"""Build a saver hydration plan from channel specs."""
|
||||
channels = tuple(
|
||||
IncrementalChannelSpec(name=name, kind=channel.checkpoint_hydration_kind)
|
||||
for name, channel in specs.items()
|
||||
if isinstance(channel, BaseChannel)
|
||||
and channel.checkpoint_hydration_kind is not None
|
||||
)
|
||||
return CheckpointHydrationPlan(channels=channels) if channels else None
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
@@ -37,82 +48,35 @@ 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.
|
||||
"""
|
||||
"""Create a checkpoint for the given channels."""
|
||||
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
|
||||
v = channels[k].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):
|
||||
return False
|
||||
return stored is MISSING or stored is DELTA_SENTINEL
|
||||
|
||||
|
||||
def channels_from_checkpoint(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
checkpoint: Checkpoint,
|
||||
*,
|
||||
saver: BaseCheckpointSaver | None = None,
|
||||
config: RunnableConfig | None = None,
|
||||
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
|
||||
"""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`.
|
||||
"""
|
||||
"""Get channels from a checkpoint."""
|
||||
channel_specs: dict[str, BaseChannel] = {}
|
||||
managed_specs: dict[str, ManagedValueSpec] = {}
|
||||
for k, v in specs.items():
|
||||
@@ -120,51 +84,10 @@ def channels_from_checkpoint(
|
||||
channel_specs[k] = v
|
||||
else:
|
||||
managed_specs[k] = v
|
||||
|
||||
channels: dict[str, BaseChannel] = {}
|
||||
for k, spec in channel_specs.items():
|
||||
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 = saver._get_channel_writes_history(config, k)
|
||||
replay_ch = spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
ch = replay_ch
|
||||
else:
|
||||
ch = spec.from_checkpoint(stored)
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
|
||||
|
||||
async def achannels_from_checkpoint(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
checkpoint: Checkpoint,
|
||||
*,
|
||||
saver: BaseCheckpointSaver | None = None,
|
||||
config: RunnableConfig | None = None,
|
||||
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
|
||||
"""Async version of `channels_from_checkpoint`. See docstring there."""
|
||||
channel_specs: dict[str, BaseChannel] = {}
|
||||
managed_specs: dict[str, ManagedValueSpec] = {}
|
||||
for k, v in specs.items():
|
||||
if isinstance(v, BaseChannel):
|
||||
channel_specs[k] = v
|
||||
else:
|
||||
managed_specs[k] = v
|
||||
|
||||
channels: dict[str, BaseChannel] = {}
|
||||
for k, spec in channel_specs.items():
|
||||
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)
|
||||
ch = replay_ch
|
||||
else:
|
||||
ch = spec.from_checkpoint(stored)
|
||||
for k, v in channel_specs.items():
|
||||
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
ch.after_checkpoint(checkpoint["channel_versions"].get(k), checkpoint.get("id"))
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from contextlib import (
|
||||
ExitStack,
|
||||
)
|
||||
from datetime import datetime, timezone
|
||||
from functools import cached_property
|
||||
from inspect import signature
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
@@ -25,11 +26,11 @@ 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,
|
||||
Checkpoint,
|
||||
CheckpointHydrationPlan,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
PendingWrite,
|
||||
@@ -93,8 +94,8 @@ from langgraph.pregel._algo import (
|
||||
task_path_str,
|
||||
)
|
||||
from langgraph.pregel._checkpoint import (
|
||||
achannels_from_checkpoint,
|
||||
channels_from_checkpoint,
|
||||
checkpoint_hydration_plan,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
@@ -199,7 +200,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",
|
||||
@@ -317,6 +317,29 @@ class PregelLoop:
|
||||
)
|
||||
self.prev_checkpoint_config = None
|
||||
|
||||
@cached_property
|
||||
def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None:
|
||||
"""Build the saver hydration plan from this loop's channel specs."""
|
||||
return checkpoint_hydration_plan(self.specs)
|
||||
|
||||
def _materialize_saved_checkpoint(
|
||||
self, saved: CheckpointTuple | None
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or self.checkpointer is None:
|
||||
return saved
|
||||
return self.checkpointer.materialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
async def _amaterialize_saved_checkpoint(
|
||||
self, saved: CheckpointTuple | None
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or self.checkpointer is None:
|
||||
return saved
|
||||
return await self.checkpointer.amaterialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
def _push_graph_lifecycle_event(
|
||||
self,
|
||||
kind: Literal["resume", "interrupt"],
|
||||
@@ -409,7 +432,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 +440,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)
|
||||
@@ -835,18 +857,8 @@ class PregelLoop:
|
||||
# parent. For forks (source=update/fork), use the fork's parent
|
||||
# checkpoint ID since the fork was created after the subgraph's
|
||||
# checkpoints from the original execution.
|
||||
#
|
||||
# Only gate on is_time_traveling (not is_replaying). When the
|
||||
# client resumes with an explicit checkpoint_id that happens to
|
||||
# point at the current head (e.g. LangGraph Studio sending
|
||||
# `checkpoint: {checkpoint_id}` alongside Command(resume=...)),
|
||||
# is_replaying is True but is_time_traveling is False. In that
|
||||
# case subgraphs should load their latest checkpoint normally,
|
||||
# not go through ReplayState's before-bound lookup which would
|
||||
# miss subgraph checkpoints created during processing of the
|
||||
# current parent step.
|
||||
replay_state: ReplayState | None = None
|
||||
if is_time_traveling:
|
||||
if self.is_replaying:
|
||||
replay_checkpoint_id = self.checkpoint["id"]
|
||||
if (
|
||||
self.checkpoint_metadata.get("source")
|
||||
@@ -894,10 +906,13 @@ 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,
|
||||
)
|
||||
if do_checkpoint and self.channels:
|
||||
for k, ch in self.channels.items():
|
||||
ch.after_checkpoint(
|
||||
self.checkpoint["channel_versions"].get(k),
|
||||
self.checkpoint.get("id"),
|
||||
)
|
||||
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
|
||||
if TASKS in self.checkpoint["channel_values"] and any(
|
||||
isinstance(channel, UntrackedValue) for channel in self.channels.values()
|
||||
@@ -934,17 +949,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
|
||||
@@ -1265,6 +1269,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
|
||||
saved = self._materialize_saved_checkpoint(saved)
|
||||
|
||||
if saved is None:
|
||||
saved = CheckpointTuple(
|
||||
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
|
||||
@@ -1289,13 +1295,9 @@ 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,
|
||||
self.checkpoint,
|
||||
saver=self.checkpointer,
|
||||
config=self.checkpoint_config,
|
||||
self.specs, self.checkpoint
|
||||
)
|
||||
self.stack.push(self._suppress_interrupt)
|
||||
self.status = "input"
|
||||
@@ -1471,6 +1473,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
|
||||
saved = await self._amaterialize_saved_checkpoint(saved)
|
||||
|
||||
if saved is None:
|
||||
saved = CheckpointTuple(
|
||||
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
|
||||
@@ -1495,15 +1499,11 @@ 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)
|
||||
)
|
||||
self.channels, self.managed = await achannels_from_checkpoint(
|
||||
self.specs,
|
||||
self.checkpoint,
|
||||
saver=self.checkpointer,
|
||||
config=self.checkpoint_config,
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs, self.checkpoint
|
||||
)
|
||||
self.stack.push(self._suppress_interrupt)
|
||||
self.status = "input"
|
||||
|
||||
@@ -14,7 +14,7 @@ from langchain_core.messages import BaseMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph._internal._constants import NS_END, NS_SEP
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
from langgraph.types import Command
|
||||
@@ -132,15 +132,23 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
task_checkpoint_ns = cast(str, metadata["langgraph_checkpoint_ns"])
|
||||
checkpoint_ns = (
|
||||
f"{task_checkpoint_ns.rsplit(NS_END, 1)[0]}{NS_END}"
|
||||
if NS_END in task_checkpoint_ns
|
||||
else task_checkpoint_ns
|
||||
)
|
||||
ns = tuple(task_checkpoint_ns.split(NS_SEP))[:-1]
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
stream_metadata = dict(metadata)
|
||||
stream_metadata["langgraph_checkpoint_ns"] = checkpoint_ns
|
||||
# Preserve backwards-compatible streamed checkpoint metadata shape.
|
||||
stream_metadata["checkpoint_ns"] = checkpoint_ns
|
||||
if tags:
|
||||
if filtered_tags := [t for t in tags if not t.startswith("seq:step")]:
|
||||
metadata["tags"] = filtered_tags
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
stream_metadata["tags"] = filtered_tags
|
||||
self.metadata[run_id] = (ns, stream_metadata)
|
||||
|
||||
def on_llm_new_token(
|
||||
self,
|
||||
|
||||
@@ -17,7 +17,7 @@ from collections.abc import (
|
||||
Sequence,
|
||||
)
|
||||
from dataclasses import is_dataclass, replace
|
||||
from functools import partial
|
||||
from functools import cached_property, partial
|
||||
from inspect import isclass
|
||||
from typing import (
|
||||
Any,
|
||||
@@ -44,6 +44,7 @@ from langgraph.cache.base import BaseCache
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
CheckpointHydrationPlan,
|
||||
CheckpointTuple,
|
||||
)
|
||||
from langgraph.store.base import BaseStore
|
||||
@@ -122,8 +123,8 @@ from langgraph.pregel._algo import (
|
||||
)
|
||||
from langgraph.pregel._call import identifier
|
||||
from langgraph.pregel._checkpoint import (
|
||||
achannels_from_checkpoint,
|
||||
channels_from_checkpoint,
|
||||
checkpoint_hydration_plan,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
@@ -729,6 +730,54 @@ class Pregel(
|
||||
return checkpointer
|
||||
return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist)
|
||||
|
||||
@cached_property
|
||||
def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None:
|
||||
return checkpoint_hydration_plan(self.channels)
|
||||
|
||||
def _materialize_saved_checkpoint(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: CheckpointTuple | None,
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or checkpointer is None:
|
||||
return saved
|
||||
return checkpointer.materialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
async def _amaterialize_saved_checkpoint(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: CheckpointTuple | None,
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or checkpointer is None:
|
||||
return saved
|
||||
return await checkpointer.amaterialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
def _materialize_saved_checkpoints(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: Sequence[CheckpointTuple],
|
||||
) -> list[CheckpointTuple]:
|
||||
if checkpointer is None or not saved:
|
||||
return list(saved)
|
||||
return checkpointer.materialize_checkpoint_tuples(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
async def _amaterialize_saved_checkpoints(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: Sequence[CheckpointTuple],
|
||||
) -> list[CheckpointTuple]:
|
||||
if checkpointer is None or not saved:
|
||||
return list(saved)
|
||||
return await checkpointer.amaterialize_checkpoint_tuples(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
def get_graph(
|
||||
self, config: RunnableConfig | None = None, *, xray: int | bool = False
|
||||
) -> Graph:
|
||||
@@ -1050,17 +1099,14 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
saver=self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config,
|
||||
checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1173,17 +1219,14 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
channels, managed = await achannels_from_checkpoint(
|
||||
checkpoint = saved.checkpoint
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
saver=self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config,
|
||||
checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1309,7 +1352,9 @@ class Pregel(
|
||||
if not isinstance(thread_id, str):
|
||||
config[CONF][CONFIG_KEY_THREAD_ID] = str(thread_id)
|
||||
|
||||
saved = checkpointer.get_tuple(config)
|
||||
saved = self._materialize_saved_checkpoint(
|
||||
checkpointer, checkpointer.get_tuple(config)
|
||||
)
|
||||
return self._prepare_state_snapshot(
|
||||
config,
|
||||
saved,
|
||||
@@ -1353,7 +1398,9 @@ class Pregel(
|
||||
if not isinstance(thread_id, str):
|
||||
config[CONF][CONFIG_KEY_THREAD_ID] = str(thread_id)
|
||||
|
||||
saved = await checkpointer.aget_tuple(config)
|
||||
saved = await self._amaterialize_saved_checkpoint(
|
||||
checkpointer, await checkpointer.aget_tuple(config)
|
||||
)
|
||||
return await self._aprepare_state_snapshot(
|
||||
config,
|
||||
saved,
|
||||
@@ -1407,9 +1454,11 @@ class Pregel(
|
||||
},
|
||||
)
|
||||
# eagerly consume list() to avoid holding up the db cursor
|
||||
for checkpoint_tuple in list(
|
||||
checkpointer.list(config, before=before, limit=limit, filter=filter)
|
||||
):
|
||||
checkpoint_tuples = self._materialize_saved_checkpoints(
|
||||
checkpointer,
|
||||
list(checkpointer.list(config, before=before, limit=limit, filter=filter)),
|
||||
)
|
||||
for checkpoint_tuple in checkpoint_tuples:
|
||||
yield self._prepare_state_snapshot(
|
||||
checkpoint_tuple.config, checkpoint_tuple
|
||||
)
|
||||
@@ -1461,12 +1510,16 @@ class Pregel(
|
||||
},
|
||||
)
|
||||
# eagerly consume list() to avoid holding up the db cursor
|
||||
for checkpoint_tuple in [
|
||||
c
|
||||
async for c in checkpointer.alist(
|
||||
config, before=before, limit=limit, filter=filter
|
||||
)
|
||||
]:
|
||||
checkpoint_tuples = await self._amaterialize_saved_checkpoints(
|
||||
checkpointer,
|
||||
[
|
||||
c
|
||||
async for c in checkpointer.alist(
|
||||
config, before=before, limit=limit, filter=filter
|
||||
)
|
||||
],
|
||||
)
|
||||
for checkpoint_tuple in checkpoint_tuples:
|
||||
yield await self._aprepare_state_snapshot(
|
||||
checkpoint_tuple.config, checkpoint_tuple
|
||||
)
|
||||
@@ -1526,12 +1579,13 @@ class Pregel(
|
||||
) -> RunnableConfig:
|
||||
# get last checkpoint
|
||||
config = ensure_config(self.config, input_config)
|
||||
saved = checkpointer.get_tuple(config)
|
||||
saved = self._materialize_saved_checkpoint(
|
||||
checkpointer, checkpointer.get_tuple(config)
|
||||
)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
@@ -1550,11 +1604,6 @@ class Pregel(
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saver=self.checkpointer
|
||||
if saved is not None
|
||||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config if saved is not None else None,
|
||||
)
|
||||
values, as_node = updates[0][:2]
|
||||
|
||||
@@ -1977,12 +2026,13 @@ class Pregel(
|
||||
) -> RunnableConfig:
|
||||
# get last checkpoint
|
||||
config = ensure_config(self.config, input_config)
|
||||
saved = await checkpointer.aget_tuple(config)
|
||||
saved = await self._amaterialize_saved_checkpoint(
|
||||
checkpointer, await checkpointer.aget_tuple(config)
|
||||
)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
@@ -1998,14 +2048,9 @@ class Pregel(
|
||||
)
|
||||
if saved:
|
||||
checkpoint_config = patch_configurable(config, saved.config[CONF])
|
||||
channels, managed = await achannels_from_checkpoint(
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saver=self.checkpointer
|
||||
if saved is not None
|
||||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config if saved is not None else None,
|
||||
)
|
||||
values, as_node = updates[0][:2]
|
||||
# no values, just clear all tasks
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.1.9"
|
||||
version = "1.1.7a2"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
@@ -24,7 +24,7 @@ classifiers = [
|
||||
'Programming Language :: Python :: 3.13',
|
||||
]
|
||||
dependencies = [
|
||||
"langchain-core>=1.3.0,<2",
|
||||
"langchain-core==1.3.0a2",
|
||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||
"langgraph-sdk>=0.3.0,<0.4.0",
|
||||
"langgraph-prebuilt>=1.0.9,<1.1.0",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,27 +1,21 @@
|
||||
"""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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
@@ -31,12 +25,13 @@ from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
|
||||
_POSTGRES_AVAILABLE = True
|
||||
_POSTGRES_URI = "postgres://sydney_runkle@localhost:5441/postgres?sslmode=disable"
|
||||
_SQLITE_AVAILABLE = True
|
||||
except ImportError:
|
||||
_POSTGRES_AVAILABLE = False
|
||||
_SQLITE_AVAILABLE = False
|
||||
|
||||
SNAPSHOT_EVERY = 50
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Realistic message payload (~100 tokens / ~400 chars each)
|
||||
@@ -45,9 +40,7 @@ except ImportError:
|
||||
_HUMAN_TEMPLATE = (
|
||||
"I need help understanding the implications of {topic} on our system architecture. "
|
||||
"Specifically, I'm concerned about how this interacts with our existing {concern} "
|
||||
"and whether we need to refactor the {component} layer before proceeding. "
|
||||
"We've had prior incidents in this area and want to be deliberate. "
|
||||
"What should we prioritize first, and are there known failure modes we should design around from the start?"
|
||||
"and whether we need to refactor the {component} layer before proceeding."
|
||||
)
|
||||
|
||||
_AI_TEMPLATE = (
|
||||
@@ -117,16 +110,8 @@ 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]},
|
||||
)
|
||||
class DeltaSnapshotState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=SNAPSHOT_EVERY)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -172,8 +157,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 +172,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 +194,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,248 +204,146 @@ def _approx_tokens(n_turns: int) -> str:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpointer factories
|
||||
# Benchmark matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@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,))
|
||||
# Turn counts chosen to span from a short session to a long-running agent conversation.
|
||||
# Storage and time complexity differences are clearly visible by 500 turns.
|
||||
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
|
||||
TURN_COUNTS = [50, 100, 200, 500]
|
||||
|
||||
|
||||
def _checkpointers() -> list[tuple[str, Any]]:
|
||||
"""Return (label, saver_or_None) pairs for available checkpointers."""
|
||||
result: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
if _POSTGRES_AVAILABLE:
|
||||
try:
|
||||
import psycopg
|
||||
def _checkpointer_factories() -> list[tuple[str, Any]]:
|
||||
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
|
||||
factories: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
if _SQLITE_AVAILABLE:
|
||||
import tempfile
|
||||
|
||||
psycopg.connect(_POSTGRES_URI).close()
|
||||
result.append(("Postgres", "postgres"))
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
factories.append(("SQLite", tempfile.NamedTemporaryFile(suffix=".db")))
|
||||
return factories
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part 1: baseline DeltaChannel(inf) vs add_messages
|
||||
# ---------------------------------------------------------------------------
|
||||
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()
|
||||
|
||||
BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
DELTA_ONLY_TURN_COUNTS = [1000]
|
||||
checkpointers: list[tuple[str, Any]] = [("InMemory (fast-path)", None)]
|
||||
if _SQLITE_AVAILABLE:
|
||||
checkpointers.append(("SQLite (get_tuple fallback)", "sqlite"))
|
||||
|
||||
for cp_label, cp_hint in checkpointers:
|
||||
print(f"--- Checkpointer: {cp_label} ---")
|
||||
_run_benchmark_for_checkpointer(cp_hint)
|
||||
|
||||
|
||||
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
|
||||
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:
|
||||
W = 120
|
||||
print("=" * W)
|
||||
header = (
|
||||
f"{'turns':>6} {'ctx size':>10} "
|
||||
f"{'add_msgs (bytes)':>18} {'delta (bytes)':>15} {'delta+snap (bytes)':>18} "
|
||||
f"{'storage saved':>14} "
|
||||
f"{'read: add_msgs':>14} {'read: delta+snap':>16}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * W)
|
||||
|
||||
results = []
|
||||
for turns in TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt))
|
||||
for turns in DELTA_ONLY_TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, None, d_bytes, None, d_rt, None, d_wt))
|
||||
s_wt, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
|
||||
|
||||
def _bytes_or_na(v: Any) -> str:
|
||||
if v is None or v < 0:
|
||||
return "n/a"
|
||||
return _fmt_bytes(v)
|
||||
|
||||
def _ms_or_na(v: Any) -> str:
|
||||
return "n/a" if v is None else f"{v * 1000:.1f}ms"
|
||||
|
||||
print(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"
|
||||
# For non-InMemory savers, blob_bytes are unavailable (-1); use read times only
|
||||
if b_bytes < 0 or s_bytes < 0:
|
||||
b_bytes_str = "n/a"
|
||||
d_bytes_str = "n/a"
|
||||
s_bytes_str = "n/a"
|
||||
storage_ratio_str = "n/a"
|
||||
else:
|
||||
ratio = b_bytes / d_bytes if d_bytes else float("inf")
|
||||
ratio_str = f"{ratio:.0f}x"
|
||||
storage_ratio = b_bytes / s_bytes if s_bytes else float("inf")
|
||||
b_bytes_str = _fmt_bytes(b_bytes)
|
||||
d_bytes_str = _fmt_bytes(d_bytes)
|
||||
s_bytes_str = _fmt_bytes(s_bytes)
|
||||
storage_ratio_str = f"{storage_ratio:.1f}x"
|
||||
results.append((turns, b_bytes, s_bytes, b_rt, s_rt, storage_ratio))
|
||||
|
||||
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"{b_bytes_str:>18} {d_bytes_str:>15} {s_bytes_str:>18} "
|
||||
f"{storage_ratio_str:>14} "
|
||||
f"{b_rt * 1000:>12.1f}ms {s_rt * 1000:>14.1f}ms"
|
||||
)
|
||||
|
||||
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))
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
if results:
|
||||
best = results[-1]
|
||||
turns, b_bytes, s_bytes, b_rt, s_rt, ratio = best
|
||||
print(f"Key findings at max scale ({turns} turns):")
|
||||
print(
|
||||
f" {turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
|
||||
f" Storage: {_fmt_bytes(b_bytes)} (add_messages) → {_fmt_bytes(s_bytes)} (DeltaChannel+snapshot) — {ratio:.0f}x reduction"
|
||||
)
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part 2: snapshot_frequency sweep
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Frequencies to test. 1 = always snapshot (like BinOp), inf = pure delta.
|
||||
SNAPSHOT_FREQUENCIES: list[int | float] = [1, 5, 10, 50, math.inf]
|
||||
|
||||
# Turn counts for the sweep — high enough to show storage divergence.
|
||||
SWEEP_TURN_COUNTS = [50, 100, 500]
|
||||
|
||||
|
||||
def _freq_label(freq: int | float) -> str:
|
||||
if freq == math.inf:
|
||||
return "inf"
|
||||
return str(int(freq))
|
||||
|
||||
|
||||
def _run_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()
|
||||
print(
|
||||
f" Read latency: {b_rt * 1000:.1f}ms (add_messages) vs {s_rt * 1000:.1f}ms (DeltaChannel+snapshot)"
|
||||
)
|
||||
print()
|
||||
print("Legend:")
|
||||
print(
|
||||
" freq=1 snapshot every write (full blob always — same as add_messages / BinOp)"
|
||||
" add_msgs = Annotated[list, add_messages] — current default, O(N²) storage"
|
||||
)
|
||||
print(
|
||||
" delta = DeltaChannel(add_messages) — O(N) storage, unbounded chain at read"
|
||||
)
|
||||
print(
|
||||
f" delta+snap = DeltaChannel(add_messages, snapshot_every={SNAPSHOT_EVERY}) — O(N) storage, O(1) read depth"
|
||||
)
|
||||
print(" freq=N snapshot every N writes; read walks at most N ancestor writes")
|
||||
print(" freq=inf pure delta; read walks entire ancestor chain")
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry points
|
||||
# Pytest 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()
|
||||
|
||||
for turns in [25, 50]:
|
||||
# Correctness assertion: DeltaChannel must use less storage at scale.
|
||||
for turns in [100, 200]:
|
||||
_, _, b_bytes = _run_turns(turns, BinaryState)
|
||||
_, _, d_bytes = _run_turns(turns, DeltaState)
|
||||
_, _, s_bytes = _run_turns(turns, DeltaSnapshotState)
|
||||
assert d_bytes < b_bytes, (
|
||||
f"DeltaChannel should use less storage at {turns} turns, "
|
||||
f"got delta={d_bytes} binary={b_bytes}"
|
||||
)
|
||||
|
||||
|
||||
@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"
|
||||
assert s_bytes < b_bytes, (
|
||||
f"DeltaChannel+snapshot should use less storage at {turns} turns, "
|
||||
f"got snapshot={s_bytes} binary={b_bytes}"
|
||||
)
|
||||
|
||||
|
||||
@@ -467,6 +352,5 @@ def test_snapshot_freq_benchmark(capsys: Any) -> None:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_baseline_benchmark()
|
||||
run_snapshot_freq_benchmark()
|
||||
run_benchmark()
|
||||
sys.exit(0)
|
||||
|
||||
@@ -1,504 +0,0 @@
|
||||
"""Tests for the BinaryOperatorAggregate -> DeltaChannel migration path.
|
||||
|
||||
A thread written under `BinaryOperatorAggregate(...)` must keep working
|
||||
after its annotation is swapped to `DeltaChannel(...)` on the same
|
||||
checkpointer — pre-migration state visible at each *settled* ancestor
|
||||
checkpoint is preserved, and post-migration writes fold on top through
|
||||
the reducer.
|
||||
|
||||
Mechanism under test: the saver's `_get_channel_writes_history(config,
|
||||
channel)` walks the parent chain; when it encounters an ancestor whose
|
||||
`channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it
|
||||
returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses
|
||||
it as the base value, and `replay_writes(writes)` folds on-path deltas.
|
||||
|
||||
Scenarios covered:
|
||||
|
||||
1. **Basic migration (sync + async)**: build pre-migration state with
|
||||
`BinaryOperatorAggregate`, swap the annotation to `DeltaChannel` on
|
||||
the same checkpointer, and verify that every settled pre-migration
|
||||
super-step boundary (`next=('__start__',)`) round-trips exactly
|
||||
under the delta-channel view.
|
||||
2. **Time travel into a pre-migration checkpoint** after migration —
|
||||
`graph.get_state(pre_migration_config)` at a settled ancestor
|
||||
returns the same state as under the binop channel.
|
||||
3. **Continuing a migrated thread**: driving one more super-step after
|
||||
migration produces a state that includes the pre-migration settled
|
||||
prefix plus the new delta write — proving `from_checkpoint(seed)` +
|
||||
`replay_writes` correctly fold post-migration deltas onto the
|
||||
pre-migration seed.
|
||||
4. **Base-saver fallback path**: a third-party-style subclass that
|
||||
removes the optimized `InMemorySaver` override and falls back to
|
||||
`BaseCheckpointSaver._get_channel_writes_history` must produce the
|
||||
same result as the optimized path.
|
||||
5. **Channel-type isolation across threads**: two threads on the same
|
||||
checkpointer under the delta-channel graph — one freshly-started,
|
||||
one migrated from pre-migration state — don't cross-contaminate.
|
||||
The parent-chain walk is scoped to the thread.
|
||||
|
||||
TODO: add postgres variants in the existing `libs/checkpoint-postgres`
|
||||
test files (different fixture setup; not this file).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
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
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factories
|
||||
#
|
||||
# A minimal reducer (`operator.add` on lists of str) with a noop node keeps
|
||||
# state change localized to the HumanMessage-like payload passed through
|
||||
# `invoke`. That isolates the pre/post-migration parity assertions to
|
||||
# channel-hydration semantics.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _noop(_state: Any) -> dict:
|
||||
return {}
|
||||
|
||||
|
||||
def _binop_graph(checkpointer: Any) -> Any:
|
||||
class BinopState(TypedDict):
|
||||
items: Annotated[list, BinaryOperatorAggregate(list, operator.add)]
|
||||
|
||||
return (
|
||||
StateGraph(BinopState)
|
||||
.add_node("noop", _noop)
|
||||
.add_edge(START, "noop")
|
||||
.add_edge("noop", END)
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
|
||||
|
||||
def _delta_graph(checkpointer: Any) -> Any:
|
||||
class DeltaState(TypedDict):
|
||||
items: Annotated[list, DeltaChannel(operator.add)]
|
||||
|
||||
return (
|
||||
StateGraph(DeltaState)
|
||||
.add_node("noop", _noop)
|
||||
.add_edge(START, "noop")
|
||||
.add_edge("noop", END)
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
|
||||
|
||||
def _drive(graph: Any, config: dict, tag: str, n: int) -> None:
|
||||
for i in range(n):
|
||||
graph.invoke({"items": [f"{tag}{i}"]}, config)
|
||||
|
||||
|
||||
async def _adrive(graph: Any, config: dict, tag: str, n: int) -> None:
|
||||
for i in range(n):
|
||||
await graph.ainvoke({"items": [f"{tag}{i}"]}, config)
|
||||
|
||||
|
||||
def _settled_boundaries(history: list) -> list[tuple[dict, list]]:
|
||||
"""Return `[(config, items), ...]` for every checkpoint in `history`
|
||||
whose `next == ('__start__',)` — the stable boundaries between invokes.
|
||||
"""
|
||||
return [
|
||||
(s.config, list(s.values.get("items", [])))
|
||||
for s in history
|
||||
if s.next == ("__start__",)
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Basic migration (sync + async)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_basic_migration_preserves_pre_migration_state() -> None:
|
||||
"""Build state under `BinaryOperatorAggregate`, migrate to
|
||||
`DeltaChannel` on the same checkpointer, and verify that every
|
||||
settled pre-migration super-step boundary round-trips exactly.
|
||||
|
||||
Settled boundaries (`next=('__start__',)`) are the stable hydration
|
||||
targets for the migration path: writes that produced the NEXT
|
||||
super-step are kept as `pending_writes` on the ancestor, so walking
|
||||
from a descendant finds the ancestor's blob as the seed and
|
||||
reconstructs the correct state.
|
||||
"""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "basic-sync"}}
|
||||
|
||||
# Pre-migration: accumulate items across 3 invokes.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 3)
|
||||
|
||||
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
|
||||
assert len(pre_boundaries) >= 2, "expected multiple settled boundaries"
|
||||
|
||||
# Migrate: swap the annotation on the same checkpointer.
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
for cfg, items in pre_boundaries:
|
||||
snap = delta.get_state(cfg)
|
||||
assert list(snap.values.get("items", [])) == items, (
|
||||
f"snapshot mismatch at {cfg['configurable']['checkpoint_id']}: "
|
||||
f"expected {items}, got {snap.values.get('items', [])}"
|
||||
)
|
||||
|
||||
|
||||
async def test_basic_migration_preserves_pre_migration_state_async() -> None:
|
||||
"""Async variant of the basic migration scenario."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "basic-async"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
await _adrive(binop, config, "u", 3)
|
||||
|
||||
pre_history = [s async for s in binop.aget_state_history(config)]
|
||||
pre_boundaries = _settled_boundaries(pre_history)
|
||||
assert len(pre_boundaries) >= 2
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
for cfg, items in pre_boundaries:
|
||||
snap = await delta.aget_state(cfg)
|
||||
assert list(snap.values.get("items", [])) == items, (
|
||||
f"async snapshot mismatch at {cfg['configurable']['checkpoint_id']}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Time travel into a pre-migration checkpoint after migration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_time_travel_into_pre_migration_checkpoint() -> None:
|
||||
"""After migration, `graph.get_state(pre_migration_config)` at a
|
||||
settled ancestor returns the state as stored at that point."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "time-travel"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 3)
|
||||
|
||||
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
|
||||
assert pre_boundaries, "no settled ancestors to time-travel to"
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
# Pick the oldest non-empty boundary — a long distance to walk back.
|
||||
non_empty = [(cfg, items) for cfg, items in pre_boundaries if items]
|
||||
assert non_empty, "expected at least one non-empty boundary"
|
||||
target_cfg, expected_items = non_empty[-1]
|
||||
|
||||
snap = delta.get_state(target_cfg)
|
||||
assert list(snap.values.get("items", [])) == expected_items
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Continuing a migrated thread: deltas fold onto pre-migration seed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_continuing_migrated_thread_folds_deltas_on_seed() -> None:
|
||||
"""Resume a pre-migration settled ancestor via `invoke(None, cfg)`
|
||||
under the delta-channel graph. Since the pre-migration checkpoint
|
||||
has an existing `pending_writes` entry (the input for the NEXT
|
||||
super-step), re-running from that ancestor reproduces the same
|
||||
post-ancestor state as the original binop run.
|
||||
|
||||
This proves the seed-terminator + write-replay pipeline works
|
||||
end-to-end across the migration boundary.
|
||||
"""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "continue"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 2)
|
||||
|
||||
# Pick the oldest settled boundary with non-empty state.
|
||||
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
|
||||
target_cfg, seed_items = next(
|
||||
(cfg, items) for cfg, items in reversed(pre_boundaries) if items
|
||||
)
|
||||
assert seed_items, "need a non-empty seed boundary"
|
||||
|
||||
# Migrate and resume from the pre-migration ancestor. `invoke(None,
|
||||
# cfg)` replays the pending writes staged at `cfg` under the new
|
||||
# channel; the reducer folds those deltas onto the seed.
|
||||
delta = _delta_graph(checkpointer)
|
||||
result = delta.invoke(None, target_cfg)
|
||||
|
||||
# The resumed state must include the pre-migration seed items in order.
|
||||
result_items = list(result.get("items", []))
|
||||
for idx, prefix_item in enumerate(seed_items):
|
||||
assert result_items[idx] == prefix_item, (
|
||||
f"pre-migration seed item at {idx} not preserved: "
|
||||
f"got {result_items[: idx + 1]}, expected {seed_items}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Base-saver fallback path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ThirdPartyStyleSaver(InMemorySaver):
|
||||
"""Simulates a third-party saver that inherits the reference
|
||||
`_get_channel_writes_history` implementation from
|
||||
`BaseCheckpointSaver` rather than overriding it.
|
||||
|
||||
We rebind the two methods to the base-class versions (via MRO) so
|
||||
the fallback path is exercised even though the storage layer is
|
||||
still the in-memory one.
|
||||
"""
|
||||
|
||||
# MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...]
|
||||
_get_channel_writes_history = ( # type: ignore[assignment]
|
||||
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
_aget_channel_writes_history = ( # type: ignore[assignment]
|
||||
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
|
||||
def test_base_saver_fallback_matches_optimized_override() -> None:
|
||||
"""The reference `BaseCheckpointSaver` implementation must produce
|
||||
the same migration behavior as the optimized `InMemorySaver`
|
||||
override. We drive the same migration scenario through both savers
|
||||
and assert per-snapshot parity in the delta-channel view."""
|
||||
|
||||
# Fast path: optimized InMemorySaver override.
|
||||
fast_saver = InMemorySaver()
|
||||
fast_config = {"configurable": {"thread_id": "fast"}}
|
||||
fast_binop = _binop_graph(fast_saver)
|
||||
_drive(fast_binop, fast_config, "u", 3)
|
||||
fast_delta = _delta_graph(fast_saver)
|
||||
fast_history = [
|
||||
(s.next, list(s.values.get("items", [])))
|
||||
for s in fast_delta.get_state_history(fast_config)
|
||||
]
|
||||
|
||||
# Slow path: base-class fallback.
|
||||
slow_saver = _ThirdPartyStyleSaver()
|
||||
slow_config = {"configurable": {"thread_id": "slow"}}
|
||||
slow_binop = _binop_graph(slow_saver)
|
||||
_drive(slow_binop, slow_config, "u", 3)
|
||||
slow_delta = _delta_graph(slow_saver)
|
||||
slow_history = [
|
||||
(s.next, list(s.values.get("items", [])))
|
||||
for s in slow_delta.get_state_history(slow_config)
|
||||
]
|
||||
|
||||
assert slow_history == fast_history, (
|
||||
"base-saver fallback should match optimized-override behavior; "
|
||||
f"fast={fast_history}, slow={slow_history}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Thread isolation under mixed-generation storage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_and_migrated_threads_do_not_cross_contaminate() -> None:
|
||||
"""Two threads sharing a checkpointer — one migrated from
|
||||
pre-migration state, one freshly-started under DeltaChannel — must
|
||||
maintain independent state. The parent-chain walk in
|
||||
`_get_channel_writes_history` must be scoped to the target thread.
|
||||
"""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
migrated_cfg = {"configurable": {"thread_id": "migrated"}}
|
||||
fresh_cfg = {"configurable": {"thread_id": "fresh"}}
|
||||
|
||||
# Thread A: pre-migration build-up.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, migrated_cfg, "m", 2)
|
||||
|
||||
# Thread B: fresh delta-channel run.
|
||||
delta = _delta_graph(checkpointer)
|
||||
_drive(delta, fresh_cfg, "f", 2)
|
||||
|
||||
# Thread A: migrate and confirm its state is anchored in its own
|
||||
# thread's pre-migration history (tag 'm'), never mixing in tag 'f'.
|
||||
migrated_boundaries = _settled_boundaries(
|
||||
list(delta.get_state_history(migrated_cfg))
|
||||
)
|
||||
assert migrated_boundaries, "migrated thread has no settled boundaries"
|
||||
for _, items in migrated_boundaries:
|
||||
for it in items:
|
||||
assert it.startswith("m"), (
|
||||
f"migrated thread leaked item from other thread: {it}"
|
||||
)
|
||||
|
||||
# Thread B: settled boundaries must only contain 'f' tags.
|
||||
fresh_boundaries = _settled_boundaries(list(delta.get_state_history(fresh_cfg)))
|
||||
assert fresh_boundaries, "fresh thread has no settled boundaries"
|
||||
for _, items in fresh_boundaries:
|
||||
for it in items:
|
||||
assert it.startswith("f"), (
|
||||
f"fresh thread leaked item from migrated thread: {it}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. Tip-of-pre-migration hydration: the latest checkpoint from a binop-run
|
||||
# thread has a real accumulated value in its own `channel_values["items"]`.
|
||||
# When hydrated under the delta-channel graph via `get_state(config)` with no
|
||||
# `checkpoint_id`, the short-circuit must use that value directly instead of
|
||||
# walking ancestors (which would skip the tip's own blob).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_tip_of_pre_migration_hydrates_directly() -> None:
|
||||
"""`graph.get_state(config)` at the latest (pre-migration) checkpoint
|
||||
returns the full accumulated list stored in that checkpoint's own
|
||||
`channel_values`. The hydration must not walk ancestors past it."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "tip-sync"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 3)
|
||||
|
||||
binop_tip = binop.get_state(config)
|
||||
expected_items = list(binop_tip.values.get("items", []))
|
||||
assert expected_items == ["u0", "u1", "u2"], (
|
||||
f"sanity: pre-migration tip should accumulate all 3 items, got {expected_items}"
|
||||
)
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
snap = delta.get_state(config)
|
||||
assert list(snap.values.get("items", [])) == expected_items, (
|
||||
f"tip hydration mismatch: expected {expected_items}, "
|
||||
f"got {snap.values.get('items', [])}"
|
||||
)
|
||||
|
||||
|
||||
async def test_tip_of_pre_migration_hydrates_directly_async() -> None:
|
||||
"""Async variant of the tip-of-pre-migration hydration scenario."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "tip-async"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
await _adrive(binop, config, "u", 3)
|
||||
|
||||
binop_tip = await binop.aget_state(config)
|
||||
expected_items = list(binop_tip.values.get("items", []))
|
||||
assert expected_items == ["u0", "u1", "u2"]
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
snap = await delta.aget_state(config)
|
||||
assert list(snap.values.get("items", [])) == expected_items, (
|
||||
f"async tip hydration mismatch: expected {expected_items}, "
|
||||
f"got {snap.values.get('items', [])}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. `update_state` after migration writes a real value to the new
|
||||
# checkpoint's `channel_values` (not a sentinel). Hydration must use it
|
||||
# directly — the ancestor walk would skip this blob and return stale state.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_update_state_after_migration_uses_written_value() -> None:
|
||||
"""After migrating and running at least one post-migration super-step
|
||||
(so the thread's tip has a `DELTA_SENTINEL`), `update_state` writes a
|
||||
concrete value to a new checkpoint's `channel_values`. `get_state`
|
||||
must reflect that concrete value."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "update-state"}}
|
||||
|
||||
# Pre-migration: accumulate a little state.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 2)
|
||||
|
||||
# Migrate and run one more super-step so the tip is a post-migration
|
||||
# checkpoint with `DELTA_SENTINEL` in its own `channel_values`.
|
||||
delta = _delta_graph(checkpointer)
|
||||
delta.invoke({"items": ["post"]}, config)
|
||||
|
||||
# `update_state` writes a concrete value into a new checkpoint's blob
|
||||
# via the reducer against the hydrated prior state.
|
||||
delta.update_state(config, {"items": ["x", "y"]})
|
||||
|
||||
snap = delta.get_state(config)
|
||||
updated_items = list(snap.values.get("items", []))
|
||||
# Must include the "x","y" update; without the hydration fix, the
|
||||
# update_state-written blob would be skipped in favor of an ancestor
|
||||
# walk, and the update values would disappear.
|
||||
assert "x" in updated_items and "y" in updated_items, (
|
||||
f"update_state values missing from snapshot: {updated_items}"
|
||||
)
|
||||
# The "x","y" items should be folded onto the prior accumulated state,
|
||||
# not stand alone. This verifies the update-written blob is used
|
||||
# directly by `get_state` (no ancestor walk past it).
|
||||
assert len(updated_items) >= 4, (
|
||||
f"update_state snapshot should preserve pre-update state, got {updated_items}"
|
||||
)
|
||||
assert updated_items[-2:] == ["x", "y"], (
|
||||
f"update_state deltas should be at the tail, got {updated_items}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 8. Fork from an `update_state` checkpoint: a new run branched off the
|
||||
# update_state-produced checkpoint must see that checkpoint's concrete
|
||||
# `channel_values` as its base, with new deltas folded on top.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fork_from_update_state_checkpoint() -> None:
|
||||
"""Branching a new run from the checkpoint produced by `update_state`
|
||||
must use that checkpoint's concrete blob as the base. Additional
|
||||
deltas from the forked run fold onto it through the reducer."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "fork"}}
|
||||
|
||||
# Pre-migration build-up, then migrate and add one post-migration step.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 2)
|
||||
delta = _delta_graph(checkpointer)
|
||||
delta.invoke({"items": ["post"]}, config)
|
||||
|
||||
# Apply `update_state` and capture the returned config (references
|
||||
# the new checkpoint produced by the update).
|
||||
update_cfg = delta.update_state(config, {"items": ["x", "y"]})
|
||||
|
||||
update_snap = delta.get_state(update_cfg)
|
||||
base_items = list(update_snap.values.get("items", []))
|
||||
assert "x" in base_items and "y" in base_items, (
|
||||
f"update_state values missing from snapshot: {base_items}"
|
||||
)
|
||||
assert base_items[-2:] == ["x", "y"], (
|
||||
f"sanity: update_state deltas should be at the tail, got {base_items}"
|
||||
)
|
||||
|
||||
# Fork: invoke from the update_state checkpoint with a new delta.
|
||||
forked = delta.invoke({"items": ["fork0"]}, update_cfg)
|
||||
forked_items = list(forked.get("items", []))
|
||||
# The fork must see the update_state-written blob as its base (not
|
||||
# walk past it), and the new delta must fold on top of it.
|
||||
assert forked_items[: len(base_items)] == base_items, (
|
||||
f"fork lost update_state base: base={base_items}, forked={forked_items}"
|
||||
)
|
||||
assert forked_items[-1] == "fork0", f"fork delta not appended: {forked_items}"
|
||||
@@ -275,70 +275,3 @@ def test_graph_callbacks_accept_base_callback_manager() -> None:
|
||||
|
||||
assert "__interrupt__" in first
|
||||
assert len(graph_handler.interrupt_events) == 1
|
||||
|
||||
|
||||
def test_non_graph_handler_via_add_handler_does_not_crash() -> None:
|
||||
"""Non-GraphCallbackHandler added via add_handler should not raise.
|
||||
|
||||
Libraries like opentelemetry-instrumentation-langchain monkey-patch
|
||||
BaseCallbackManager.__init__ and inject handlers via add_handler().
|
||||
These handlers inherit from BaseCallbackHandler, not
|
||||
GraphCallbackHandler. They must be silently accepted — graph lifecycle
|
||||
events will simply not be dispatched to them.
|
||||
"""
|
||||
from langgraph.callbacks import _GraphCallbackManager
|
||||
|
||||
manager = _GraphCallbackManager()
|
||||
plain_handler = _LangChainCustomEventHandler()
|
||||
|
||||
manager.add_handler(plain_handler, inherit=True)
|
||||
assert plain_handler in manager.handlers
|
||||
|
||||
|
||||
def test_non_graph_handler_does_not_receive_lifecycle_events() -> None:
|
||||
"""Non-GraphCallbackHandler added alongside a GraphCallbackHandler
|
||||
should not interfere with lifecycle event dispatch."""
|
||||
graph = _build_interrupt_graph()
|
||||
graph_handler = _GraphEventHandler()
|
||||
plain_handler = _LangChainCustomEventHandler()
|
||||
|
||||
config = {
|
||||
"configurable": {"thread_id": "graph-callback-mixed-handlers"},
|
||||
"callbacks": [plain_handler, graph_handler],
|
||||
}
|
||||
|
||||
first = graph.invoke({"answer": None}, config)
|
||||
assert "__interrupt__" in first
|
||||
|
||||
assert len(graph_handler.interrupt_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
resumed = graph.invoke(Command(resume="done"), config)
|
||||
assert resumed == {"answer": "done"}
|
||||
assert len(graph_handler.resume_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_non_graph_handler_does_not_receive_lifecycle_events_async() -> None:
|
||||
"""Async variant: non-GraphCallbackHandler should not interfere."""
|
||||
graph = _build_interrupt_graph()
|
||||
graph_handler = _GraphEventHandler()
|
||||
plain_handler = _LangChainCustomEventHandler()
|
||||
|
||||
config = {
|
||||
"configurable": {"thread_id": "graph-callback-mixed-handlers-async"},
|
||||
"callbacks": [plain_handler, graph_handler],
|
||||
}
|
||||
|
||||
first = await graph.ainvoke({"answer": None}, config)
|
||||
assert "__interrupt__" in first
|
||||
|
||||
assert len(graph_handler.interrupt_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
resumed = await graph.ainvoke(Command(resume="done"), config)
|
||||
assert resumed == {"answer": "done"}
|
||||
assert len(graph_handler.resume_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
@@ -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
|
||||
@@ -9408,6 +9407,7 @@ async def test_delta_channel_end_to_end_inmemory() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
@@ -9449,6 +9449,7 @@ async def test_delta_channel_time_travel() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
@@ -9506,6 +9507,7 @@ async def test_delta_channel_remove_message_end_to_end() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
@@ -9552,6 +9554,7 @@ async def test_delta_channel_update_by_id_end_to_end() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
@@ -9584,90 +9587,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
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
"""Sweep snapshot_every values to find the storage vs. time-travel tradeoff.
|
||||
|
||||
Run directly: python tests/test_rehydrate_sweep.py
|
||||
Run via pytest: pytest tests/test_rehydrate_sweep.py -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
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
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
REHYDRATE_SWEEP = [5, 10, 25, 50, 100, None] # None = no rehydration (pure diff)
|
||||
TURN_COUNTS = [50, 100, 250, 500]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_state(snapshot_every: int | None) -> type:
|
||||
channel = DeltaChannel(add_messages, snapshot_every=snapshot_every)
|
||||
return TypedDict("S", {"messages": Annotated[list, channel]})
|
||||
|
||||
|
||||
def _make_graph(state_cls: type) -> Any:
|
||||
def human_node(state: Any) -> dict:
|
||||
return {}
|
||||
|
||||
def ai_node(state: Any) -> dict:
|
||||
last = state["messages"][-1]
|
||||
return {"messages": [AIMessage(content=f"reply-to-{last.id}")]}
|
||||
|
||||
g = StateGraph(state_cls)
|
||||
g.add_node("human", human_node)
|
||||
g.add_node("ai", ai_node)
|
||||
g.add_edge("human", "ai")
|
||||
g.add_edge("ai", END)
|
||||
g.set_entry_point("human")
|
||||
return g.compile(checkpointer=MemorySaver())
|
||||
|
||||
|
||||
def _total_blob_bytes(saver: MemorySaver) -> int:
|
||||
total = 0
|
||||
for (_, _, _, _), (type_tag, blob) in saver.blobs.items():
|
||||
if blob is not None:
|
||||
total += len(blob)
|
||||
return total
|
||||
|
||||
|
||||
def _measure_time_travel_ms(graph: Any, config: dict) -> float:
|
||||
"""Time how long it takes to get state at the very first checkpoint (worst case)."""
|
||||
history = list(graph.get_state_history(config))
|
||||
if not history:
|
||||
return 0.0
|
||||
oldest = history[-1]
|
||||
t0 = time.perf_counter()
|
||||
graph.get_state(oldest.config)
|
||||
return (time.perf_counter() - t0) * 1000
|
||||
|
||||
|
||||
def _run(n_turns: int, snapshot_every: int | None) -> tuple[float, int, float]:
|
||||
"""Returns (write_ms, blob_bytes, time_travel_ms)."""
|
||||
state_cls = _make_state(snapshot_every)
|
||||
graph = _make_graph(state_cls)
|
||||
saver: MemorySaver = graph.checkpointer # type: ignore[assignment]
|
||||
config = {"configurable": {"thread_id": "sweep"}}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
for i in range(n_turns):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"msg-{i}", id=f"h{i}")]}, config
|
||||
)
|
||||
write_ms = (time.perf_counter() - t0) * 1000
|
||||
|
||||
blob_bytes = _total_blob_bytes(saver)
|
||||
tt_ms = _measure_time_travel_ms(graph, config)
|
||||
return write_ms, blob_bytes, tt_ms
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ASCII sparkline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sparkline(values: list[float], width: int = 20) -> str:
|
||||
bars = " ▁▂▃▄▅▆▇█"
|
||||
lo, hi = min(values), max(values)
|
||||
span = hi - lo or 1
|
||||
chars = [bars[round((v - lo) / span * (len(bars) - 1))] for v in values]
|
||||
return "".join(chars).ljust(width)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_sweep() -> None:
|
||||
label = {v: (str(v) if v is not None else "None(∞)") for v in REHYDRATE_SWEEP}
|
||||
|
||||
print()
|
||||
print("snapshot_every sweep — storage vs time-travel cost")
|
||||
print("=" * 90)
|
||||
|
||||
for turns in TURN_COUNTS:
|
||||
print(f"\n--- {turns} turns ---")
|
||||
col_w = 12
|
||||
header = (
|
||||
f"{'snapshot_every':>18} "
|
||||
f"{'blob_bytes':>{col_w}} "
|
||||
f"{'write_ms':>{col_w}} "
|
||||
f"{'time_travel_ms':>{col_w}}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * 60)
|
||||
|
||||
tt_vals: list[float] = []
|
||||
byte_vals: list[int] = []
|
||||
write_vals: list[float] = []
|
||||
rows: list[tuple] = []
|
||||
|
||||
for rv in REHYDRATE_SWEEP:
|
||||
write_ms, blob_bytes, tt_ms = _run(turns, rv)
|
||||
rows.append((rv, blob_bytes, write_ms, tt_ms))
|
||||
byte_vals.append(blob_bytes)
|
||||
write_vals.append(write_ms)
|
||||
tt_vals.append(tt_ms)
|
||||
|
||||
for rv, blob_bytes, write_ms, tt_ms in rows:
|
||||
print(
|
||||
f"{label[rv]:>18} "
|
||||
f"{blob_bytes:>{col_w},} "
|
||||
f"{write_ms:>{col_w}.1f} "
|
||||
f"{tt_ms:>{col_w}.2f}"
|
||||
)
|
||||
|
||||
print()
|
||||
print(
|
||||
f" bytes spark: [{_sparkline(byte_vals)}] "
|
||||
f"lo={min(byte_vals):,} hi={max(byte_vals):,}"
|
||||
)
|
||||
print(
|
||||
f" time-travel spark: [{_sparkline(tt_vals)}] "
|
||||
f"lo={min(tt_vals):.2f}ms hi={max(tt_vals):.2f}ms"
|
||||
)
|
||||
print(
|
||||
f" write spark: [{_sparkline(write_vals)}] "
|
||||
f"lo={min(write_vals):.1f}ms hi={max(write_vals):.1f}ms"
|
||||
)
|
||||
|
||||
print()
|
||||
print("=" * 90)
|
||||
print(
|
||||
"snapshot_every=None means pure diff (no snapshots) — "
|
||||
"lowest storage, highest time-travel cost."
|
||||
)
|
||||
print(
|
||||
"Lower snapshot_every = more frequent full snapshots = "
|
||||
"faster time-travel, more storage."
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
def test_rehydrate_sweep(capsys: Any) -> None:
|
||||
with capsys.disabled():
|
||||
run_sweep()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_sweep()
|
||||
sys.exit(0)
|
||||
@@ -1113,70 +1113,6 @@ def test_subgraph_interrupt_replay_from_parent_then_resume(
|
||||
]
|
||||
|
||||
|
||||
def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Resume with Command(resume=...) plus the current head checkpoint_id
|
||||
in config. The subgraph must continue from the interrupted node, not
|
||||
restart from scratch. Explicit checkpoint_id triggers is_replaying but
|
||||
this is a resume, not a time-travel, so ReplayState should not apply."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["sub_a"]}
|
||||
|
||||
def ask_human(state: State) -> State:
|
||||
called.append("ask_human")
|
||||
answer = interrupt("Provide input:")
|
||||
return {"value": [f"human:{answer}"]}
|
||||
|
||||
def step_b(state: State) -> State:
|
||||
called.append("step_b")
|
||||
return {"value": ["sub_b"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_human", ask_human)
|
||||
.add_node("step_b", step_b)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_human")
|
||||
.add_edge("ask_human", "step_b")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("subgraph_node", subgraph)
|
||||
.add_edge(START, "subgraph_node")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until interrupt fires in subgraph
|
||||
graph.invoke({"value": []}, config)
|
||||
assert called == ["step_a", "ask_human"]
|
||||
|
||||
# Resume with explicit head checkpoint_id in config
|
||||
head_checkpoint_id = graph.get_state(config).config["configurable"]["checkpoint_id"]
|
||||
called.clear()
|
||||
resume_config = {
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_id": head_checkpoint_id,
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
result = graph.invoke(Command(resume="answer"), resume_config)
|
||||
|
||||
assert called == ["ask_human", "step_b"]
|
||||
assert "__interrupt__" not in result
|
||||
assert result["value"] == ["sub_a", "human:answer", "sub_b"]
|
||||
|
||||
|
||||
def test_subgraph_replay_loads_accumulated_state_then_resume(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
|
||||
Generated
+7
-7
@@ -1348,7 +1348,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.0"
|
||||
version = "1.3.0a2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -1360,14 +1360,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.9"
|
||||
version = "1.1.7a2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1439,7 +1439,7 @@ test = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
@@ -1706,7 +1706,7 @@ inmem = [
|
||||
requires-dist = [
|
||||
{ name = "click", specifier = ">=8.1.7" },
|
||||
{ name = "httpx", specifier = ">=0.24.0" },
|
||||
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.9.0" },
|
||||
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.8.0" },
|
||||
{ name = "langgraph-runtime-inmem", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.7" },
|
||||
{ name = "langgraph-sdk", marker = "python_full_version >= '3.11'", specifier = ">=0.1.0" },
|
||||
{ name = "pathspec", specifier = ">=0.11.0" },
|
||||
@@ -1742,7 +1742,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.10"
|
||||
version = "1.0.9"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -808,7 +808,6 @@ class ToolNode(RunnableCallable):
|
||||
context=runtime.context,
|
||||
store=runtime.store,
|
||||
stream_writer=runtime.stream_writer,
|
||||
tools=list(self.tools_by_name.values()),
|
||||
execution_info=runtime.execution_info,
|
||||
server_info=runtime.server_info,
|
||||
)
|
||||
@@ -843,7 +842,6 @@ class ToolNode(RunnableCallable):
|
||||
context=runtime.context,
|
||||
store=runtime.store,
|
||||
stream_writer=runtime.stream_writer,
|
||||
tools=list(self.tools_by_name.values()),
|
||||
execution_info=runtime.execution_info,
|
||||
server_info=runtime.server_info,
|
||||
)
|
||||
@@ -1578,7 +1576,6 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
|
||||
- `context`: Runtime context (shared with `Runtime`)
|
||||
- `store`: `BaseStore` instance for persistent storage (shared with `Runtime`)
|
||||
- `stream_writer`: `StreamWriter` for streaming output (shared with `Runtime`)
|
||||
- `tools`: List of all available `BaseTool` instances
|
||||
|
||||
No `Annotated` wrapper is needed - just use `runtime: ToolRuntime`
|
||||
as a parameter.
|
||||
@@ -1621,7 +1618,6 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
|
||||
context: ContextT
|
||||
config: RunnableConfig
|
||||
stream_writer: StreamWriter
|
||||
tools: list[BaseTool]
|
||||
tool_call_id: str | None
|
||||
store: BaseStore | None
|
||||
execution_info: ExecutionInfo | None = None
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.10"
|
||||
version = "1.0.9"
|
||||
description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -4,7 +4,7 @@ This tests the fix for https://github.com/langchain-ai/langchain/issues/35585
|
||||
|
||||
When using InjectedState(<field>) on a tool parameter, and the referenced field is
|
||||
declared as NotRequired in the custom state schema, the ToolNode should gracefully
|
||||
handle missing fields by injecting None instead of raising KeyError.
|
||||
handle missing fields without raising KeyError so the tool's default can apply.
|
||||
"""
|
||||
|
||||
import sys
|
||||
@@ -45,6 +45,14 @@ def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> st
|
||||
return f"It's always sunny in {city}!"
|
||||
|
||||
|
||||
@tool
|
||||
def get_weather_with_default(
|
||||
city: Annotated[str, InjectedState("city")] = "Boston",
|
||||
) -> str:
|
||||
"""Get weather for a given city, defaulting when state omits the field."""
|
||||
return f"It's always sunny in {city}!"
|
||||
|
||||
|
||||
def _create_mock_runtime(
|
||||
state: dict | None = None,
|
||||
store=None,
|
||||
@@ -69,7 +77,6 @@ def _create_config_with_runtime(store=None, state=None):
|
||||
context={},
|
||||
store=store,
|
||||
stream_writer=None,
|
||||
tools=[],
|
||||
tool_call_id="test_id",
|
||||
)
|
||||
return {
|
||||
@@ -85,7 +92,7 @@ def _create_config_with_runtime(store=None, state=None):
|
||||
reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
|
||||
)
|
||||
def test_injected_state_not_required_field_missing_injects_none():
|
||||
"""Test that InjectedState with NotRequired field injects None when field is missing.
|
||||
"""Test that missing optional InjectedState leaves the tool default in place.
|
||||
|
||||
This verifies the fix for https://github.com/langchain-ai/langchain/issues/35585
|
||||
"""
|
||||
@@ -115,6 +122,37 @@ def test_injected_state_not_required_field_missing_injects_none():
|
||||
assert "No city provided" in tool_msg.content
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
|
||||
)
|
||||
def test_injected_state_not_required_field_missing_preserves_tool_default():
|
||||
"""Test that missing optional InjectedState preserves a non-None tool default."""
|
||||
tool_node = ToolNode([get_weather_with_default])
|
||||
|
||||
tool_call = {
|
||||
"name": "get_weather_with_default",
|
||||
"args": {},
|
||||
"id": "call_1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
ai_msg = AIMessage("Let me check the weather", tool_calls=[tool_call])
|
||||
|
||||
state_without_city: CustomAgentStateWithNotRequired = {
|
||||
"messages": [HumanMessage("What's the weather?"), ai_msg],
|
||||
}
|
||||
|
||||
result = tool_node.invoke(
|
||||
state_without_city,
|
||||
config=_create_config_with_runtime(state=state_without_city),
|
||||
)
|
||||
|
||||
assert len(result["messages"]) == 1
|
||||
tool_msg = result["messages"][0]
|
||||
assert isinstance(tool_msg, ToolMessage)
|
||||
assert "Boston" in tool_msg.content
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
|
||||
|
||||
@@ -2016,8 +2016,8 @@ async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async()
|
||||
assert tool_message.tool_call_id == "call_dynamic_2"
|
||||
|
||||
|
||||
def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
"""Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""
|
||||
def test_tool_runtime_forwards_execution_info_and_server_info() -> None:
|
||||
"""Test that execution_info and server_info are forwarded from Runtime to ToolRuntime."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
|
||||
exec_info = ExecutionInfo(
|
||||
@@ -2043,15 +2043,9 @@ def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
"""Tool that captures runtime info."""
|
||||
captured["execution_info"] = runtime.execution_info
|
||||
captured["server_info"] = runtime.server_info
|
||||
captured["tools"] = runtime.tools
|
||||
return "ok"
|
||||
|
||||
@dec_tool
|
||||
def other_tool(y: int) -> str:
|
||||
"""Another tool available to the runtime."""
|
||||
return str(y)
|
||||
|
||||
node = ToolNode([info_tool, other_tool])
|
||||
node = ToolNode([info_tool])
|
||||
tool_call = {
|
||||
"name": "info_tool",
|
||||
"args": {"x": 1},
|
||||
@@ -2060,21 +2054,17 @@ def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
||||
result = node.invoke({"messages": [msg]}, config=config)
|
||||
node.invoke({"messages": [msg]}, config=config)
|
||||
|
||||
assert result["messages"][-1].content == "ok"
|
||||
assert captured["execution_info"] is exec_info
|
||||
assert captured["execution_info"].thread_id == "t-1"
|
||||
assert captured["execution_info"].task_id == "tk-1"
|
||||
assert captured["server_info"] is server_info
|
||||
assert captured["server_info"].assistant_id == "asst-1"
|
||||
assert [tool.name for tool in captured["tools"]] == ["info_tool", "other_tool"]
|
||||
|
||||
|
||||
async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async() -> (
|
||||
None
|
||||
):
|
||||
"""Test that execution_info, server_info, and tools are forwarded in async path."""
|
||||
async def test_tool_runtime_forwards_execution_info_and_server_info_async() -> None:
|
||||
"""Test that execution_info and server_info are forwarded in async path."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
|
||||
exec_info = ExecutionInfo(
|
||||
@@ -2100,15 +2090,9 @@ async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async(
|
||||
"""Async tool that captures runtime info."""
|
||||
captured["execution_info"] = runtime.execution_info
|
||||
captured["server_info"] = runtime.server_info
|
||||
captured["tools"] = runtime.tools
|
||||
return "ok"
|
||||
|
||||
@dec_tool
|
||||
async def other_tool_async(y: int) -> str:
|
||||
"""Another async tool available to the runtime."""
|
||||
return str(y)
|
||||
|
||||
node = ToolNode([info_tool_async, other_tool_async])
|
||||
node = ToolNode([info_tool_async])
|
||||
tool_call = {
|
||||
"name": "info_tool_async",
|
||||
"args": {"x": 1},
|
||||
@@ -2117,17 +2101,12 @@ async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async(
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
||||
result = await node.ainvoke({"messages": [msg]}, config=config)
|
||||
await node.ainvoke({"messages": [msg]}, config=config)
|
||||
|
||||
assert result["messages"][-1].content == "ok"
|
||||
assert captured["execution_info"] is exec_info
|
||||
assert captured["execution_info"].thread_id == "t-2"
|
||||
assert captured["server_info"] is server_info
|
||||
assert captured["server_info"].graph_id == "graph-2"
|
||||
assert [tool.name for tool in captured["tools"]] == [
|
||||
"info_tool_async",
|
||||
"other_tool_async",
|
||||
]
|
||||
|
||||
|
||||
# --- InjectedToolArg security tests ---
|
||||
|
||||
Generated
+6
-6
@@ -249,7 +249,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.0"
|
||||
version = "1.3.0a2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -261,14 +261,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.9"
|
||||
version = "1.1.7a2"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -281,7 +281,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "." },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
@@ -490,7 +490,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.10"
|
||||
version = "1.0.9"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Generated
+6
-6
@@ -262,7 +262,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.0"
|
||||
version = "1.3.0a2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -274,14 +274,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.9"
|
||||
version = "1.1.7a2"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -294,7 +294,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "." },
|
||||
@@ -413,7 +413,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.10"
|
||||
version = "1.0.9"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -1,137 +0,0 @@
|
||||
# Delta-channel reconstruction: query strategy benchmark
|
||||
|
||||
**Branch:** `delta-channel-writes-based`
|
||||
**Question (Nuno):** Is the recursive CTE the right query shape for reconstructing a delta channel inside `get_tuple`, or would a plain `SELECT WHERE` be cheaper even though it returns more rows?
|
||||
**Answer:** Plain `SELECT WHERE` wins at every realistic depth. The recursion isn't the problem — the JSON-expression join inside the CTE is.
|
||||
|
||||
## Setup
|
||||
|
||||
- Postgres 16 on `localhost:5441` (the `compose-postgres.yml` instance, run directly without docker for this round)
|
||||
- Single delta channel `messages`, one write per checkpoint, `DELTA_SENTINEL` blob per checkpoint
|
||||
- Linear chain (`branch=1`) and 5-way branching at every step (`branch=5`) — branching is the case where plain over-fetches sibling rows
|
||||
- Median of 20 timed runs after 3 warmups, fresh psycopg cursor per strategy
|
||||
- Bench script: `bench_get_tuple_strategies.py` at repo root
|
||||
|
||||
Three strategies compared:
|
||||
|
||||
| name | roundtrips | shape |
|
||||
|------|-----------|-------|
|
||||
| `cte` | 1 | Current prod: recursive CTE walks ancestors, LEFT JOINs writes + blobs |
|
||||
| `plain` | 3 | Nuno's suggestion: thread-wide `SELECT WHERE` per table, Python walks parent chain and filters |
|
||||
| `cte+narrow` | 2 | CTE returns ancestor IDs only, then one `UNION ALL` of writes + blobs filtered by `ANY(ids)` |
|
||||
|
||||
## Results (ms per get_tuple, median of 20)
|
||||
|
||||
```
|
||||
depth branch cte plain cte+narrow rows_cte rows_plain plain/cte
|
||||
10 1 0.14ms 0.26ms 0.24ms 9 30 1.89x
|
||||
10 5 0.21ms 0.23ms 0.17ms 9 110 1.11x
|
||||
50 1 0.89ms 0.27ms 0.33ms 49 150 0.31x
|
||||
50 5 2.35ms 0.66ms 0.51ms 49 550 0.28x
|
||||
200 1 11.61ms 0.78ms 1.30ms 199 600 0.07x
|
||||
200 5 34.79ms 2.33ms 3.07ms 199 2200 0.07x
|
||||
1000 1 274.60ms 2.59ms 13.29ms 999 3000 0.01x
|
||||
1000 5 856.01ms 10.14ms 15.31ms 999 11000 0.01x
|
||||
```
|
||||
|
||||
Lower is better. `plain/cte < 1` means plain is faster.
|
||||
|
||||
### Headline numbers
|
||||
|
||||
- depth 50: plain is **3x** faster
|
||||
- depth 200: plain is **15x** faster
|
||||
- depth 1000: plain is **~100x** faster
|
||||
- Branching makes plain over-fetch (3000 rows → 11000 rows at d=1000), but it remains ~85x faster than the CTE
|
||||
|
||||
## Why the CTE collapses
|
||||
|
||||
`EXPLAIN (ANALYZE, BUFFERS)` of the CTE at depth 1000 (linear). Excerpt with the load-bearing nodes:
|
||||
|
||||
```
|
||||
Sort ... actual time=137.798..137.827 rows=999
|
||||
CTE ancestors
|
||||
-> Recursive Union ... actual time=0.005..2.443 rows=999
|
||||
^^^^^^
|
||||
recursion is 2.4 ms — fine
|
||||
-> Nested Loop Left Join ... actual time=2.676..137.529 rows=999
|
||||
Join Filter: (cw.checkpoint_id = a.cid)
|
||||
Rows Removed by Join Filter: 998001
|
||||
^^^^^^^
|
||||
999 ancestors x ~1000 writes
|
||||
-> Nested Loop Left Join ... actual time=2.669..85.061 rows=999
|
||||
Join Filter: (bl.version = ((c.checkpoint -> 'channel_versions'::text) ->> bl.channel))
|
||||
Rows Removed by Join Filter: 998001
|
||||
^^^^^^^
|
||||
same quadratic blow-up on the blob join
|
||||
```
|
||||
|
||||
Two pathological things are happening:
|
||||
|
||||
1. **The blob join filter is on a JSON expression**: `bl.version = (c.checkpoint -> 'channel_versions' ->> bl.channel)`. The planner cannot push this into an index lookup, so it materializes `checkpoint_blobs` for the thread and does a nested-loop comparison against every ancestor — a Cartesian product that grows as `O(ancestors × blobs_in_thread)`.
|
||||
2. **The writes join is similar**: writes for the thread are materialized once, then nested-loop joined against ancestors with a `Join Filter` rather than a hash/merge join over the indexed `checkpoint_id`.
|
||||
|
||||
At depth 1000 that's **~2 million rows evaluated, 99.9% of them discarded**. The recursion itself is a rounding error.
|
||||
|
||||
For comparison, the plain Q1 (`SELECT … FROM checkpoints WHERE thread_id=? AND checkpoint_ns=?`) at depth 1000:
|
||||
|
||||
```
|
||||
Seq Scan on checkpoints ... actual time=0.012..0.121 rows=1000
|
||||
Execution Time: 0.140 ms
|
||||
```
|
||||
|
||||
A simple seq scan over 57 buffers. Q2 and Q3 follow the same shape and complete in well under 1 ms each.
|
||||
|
||||
## Crossover and remote-DB reasoning
|
||||
|
||||
- Pure local Postgres: plain wins from depth ~30 onward; CTE wins by fractions of a ms below that
|
||||
- Remote Postgres at ~5 ms RTT adds ~10 ms to plain (3 roundtrips vs 1). Crossover shifts to ~depth 30. Above that, the CTE's quadratic SQL cost still dominates the RTT savings.
|
||||
|
||||
There is no realistic conversation depth where the CTE wins on a remote DB. At depth 200+ (anything resembling a real multi-turn agent run) plain is faster regardless of network.
|
||||
|
||||
## Recommendation
|
||||
|
||||
**Switch to plain SELECT WHERE, one delta channel at a time.**
|
||||
|
||||
Three indexed queries per delta channel:
|
||||
|
||||
```sql
|
||||
-- Q1: parent chain + per-checkpoint version of this channel
|
||||
SELECT checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> 'channel_name' AS ver
|
||||
FROM checkpoints
|
||||
WHERE thread_id = ? AND checkpoint_ns = ?;
|
||||
|
||||
-- Q2: writes for this channel, anywhere in the thread
|
||||
SELECT checkpoint_id, type, blob, task_id, idx
|
||||
FROM checkpoint_writes
|
||||
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
|
||||
|
||||
-- Q3: blobs for this channel, anywhere in the thread
|
||||
SELECT version, type, blob
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
|
||||
```
|
||||
|
||||
Python then:
|
||||
- Builds `parent_of: dict[cid, parent_cid]` from Q1
|
||||
- Walks from target's parent newest → oldest
|
||||
- Filters Q2 rows by `ancestor_set`, processes oldest → newest, applies overwrite-terminator
|
||||
- Picks seed blob via the per-ancestor `ver` map, terminates at first non-sentinel blob
|
||||
|
||||
All O(n) on n = thread checkpoints, with tight constants (dict lookups). No recursion, no JSON-expression joins, no quadratic plans.
|
||||
|
||||
If the 3-roundtrip cost ever shows up on remote-DB benchmarks, fold Q2 + Q3 into one `UNION ALL` to get back to 2 roundtrips. Bench says it isn't worth the SQL complexity right now.
|
||||
|
||||
## Bonus: code simplification from single-channel scope
|
||||
|
||||
Multi-channel reconstruction in the current `_reconstruct_delta_channels_cur` carries:
|
||||
|
||||
- `rows_by_cid` nested dicts, keyed by cid then channel
|
||||
- `seen_blob: set[(cid, channel)]` and `seen_write: set[(cid, channel, task_id, idx)]` dedup
|
||||
- `collected: dict[channel, list]`, `done: set[channel]`, `seeds: dict[channel, value]`
|
||||
- Inner `for ch in channels_list` loops and an early-exit `if len(done) == len(channels_list)`
|
||||
|
||||
Single-channel collapses these to a single list, a single bool, and one `Optional[Any]`. Roughly half the Python in that function, plus an obvious shape for splitting pure post-processing into `base.py` so sync and async stop duplicating it.
|
||||
|
||||
If multi-channel coalescing turns out to matter later, it can come back as a SQL-level optimization without re-introducing the bookkeeping in Python.
|
||||
Reference in New Issue
Block a user