mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-27 01:52:25 +02:00
Compare commits
14
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
38bf050ab5 | ||
|
|
f4b878a85a | ||
|
|
4a1d81611b | ||
|
|
a529b9bede | ||
|
|
0a26b471d3 | ||
|
|
b674dd4622 | ||
|
|
8df0a377d0 | ||
|
|
216cf33a54 | ||
|
|
4956134a37 | ||
|
|
aa94790f36 | ||
|
|
e002711ede | ||
|
|
f44b49b33d | ||
|
|
a0a95df2ac | ||
|
|
d194c18c06 |
@@ -100,4 +100,3 @@ dmypy.json
|
||||
.turbo
|
||||
.editorconfig
|
||||
.scratch
|
||||
.worktrees/
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,405 +0,0 @@
|
||||
# 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` |
|
||||
@@ -430,43 +430,6 @@ class PostgresSaver(BasePostgresSaver):
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
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 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:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
@@ -479,13 +442,6 @@ class PostgresSaver(BasePostgresSaver):
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
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": {
|
||||
@@ -498,7 +454,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**channel_values,
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
|
||||
@@ -391,43 +391,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
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 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 = 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:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
@@ -440,19 +403,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
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)
|
||||
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
@@ -460,15 +415,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**channel_values,
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -185,22 +185,15 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
)
|
||||
|
||||
def _load_blobs(
|
||||
self,
|
||||
blob_values: list[tuple[bytes, bytes, bytes]],
|
||||
*,
|
||||
thread_id: str = "",
|
||||
checkpoint_ns: str = "",
|
||||
cur: Any = None,
|
||||
self, blob_values: list[tuple[bytes, bytes, bytes]]
|
||||
) -> 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[channel] = self.serde.loads_typed((type_tag, v))
|
||||
return result
|
||||
return {
|
||||
k.decode(): self.serde.loads_typed((t.decode(), v))
|
||||
for k, t, v in blob_values
|
||||
if t.decode() != "empty"
|
||||
}
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
|
||||
@@ -371,47 +371,3 @@ async def test_get_checkpoint_no_channel_values(
|
||||
|
||||
checkpoint = await saver.aget_tuple(config)
|
||||
assert checkpoint.checkpoint["channel_values"] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
|
||||
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
|
||||
pytest.importorskip(
|
||||
"langgraph.channels.delta", reason="langgraph core not installed"
|
||||
)
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
|
||||
async with _saver(saver_name) as saver:
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "diff-channel-test-1"}}
|
||||
|
||||
await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
|
||||
await graph.ainvoke(
|
||||
{"messages": [HumanMessage(content="there", id="h2")]}, config
|
||||
)
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
msgs = state.values["messages"]
|
||||
assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].content == "hi"
|
||||
assert msgs[1].content == "reply-1"
|
||||
assert msgs[2].content == "there"
|
||||
assert msgs[3].content == "reply-3"
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import dataclasses
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
|
||||
from typing import ( # noqa: UP035
|
||||
@@ -29,26 +28,6 @@ from langgraph.checkpoint.serde.types import (
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
PendingWrite = tuple[str, str, Any]
|
||||
|
||||
|
||||
@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 by the pregel layer."""
|
||||
|
||||
base: list[Any] | None # starting accumulated value; None = start from empty
|
||||
deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -478,42 +457,6 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
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.
|
||||
"""
|
||||
return NotImplemented
|
||||
|
||||
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).
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
|
||||
@@ -126,46 +126,12 @@ class InMemorySaver(
|
||||
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":
|
||||
channel_values[k] = self.serde.loads_typed(vv)
|
||||
if kk in self.blobs:
|
||||
vv = self.blobs[kk]
|
||||
if vv[0] != "empty":
|
||||
channel_values[k] = self.serde.loads_typed(vv)
|
||||
return channel_values
|
||||
|
||||
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, {})
|
||||
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)
|
||||
|
||||
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.
|
||||
|
||||
|
||||
@@ -80,8 +80,6 @@ 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"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -46,11 +46,22 @@ 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
|
||||
|
||||
return isinstance(obj, DeltaValue)
|
||||
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)
|
||||
|
||||
|
||||
class JsonPlusSerializer(SerializerProtocol):
|
||||
@@ -245,8 +256,6 @@ 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)
|
||||
@@ -269,13 +278,6 @@ 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:
|
||||
@@ -549,7 +551,9 @@ def _create_msgpack_ext_hook(
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
logger.warning(
|
||||
_warn_once(
|
||||
_warned_unregistered_types,
|
||||
key,
|
||||
"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 "
|
||||
@@ -571,7 +575,9 @@ def _create_msgpack_ext_hook(
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
logger.warning(
|
||||
_warn_once(
|
||||
_warned_blocked_types,
|
||||
key,
|
||||
"Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
|
||||
"Add to allowed_msgpack_modules to allow: [(%r, %r)]",
|
||||
module,
|
||||
|
||||
@@ -29,6 +29,8 @@ from langgraph.checkpoint.serde.jsonplus import (
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
|
||||
|
||||
@@ -102,6 +104,13 @@ 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,6 +35,8 @@ 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
|
||||
|
||||
@@ -580,6 +582,14 @@ 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
|
||||
@@ -595,6 +605,12 @@ 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
|
||||
|
||||
|
||||
@@ -639,7 +655,6 @@ 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"))
|
||||
@@ -657,7 +672,6 @@ 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")]
|
||||
)
|
||||
@@ -983,31 +997,3 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
|
||||
# No blocking should occur - inner is serialized as dict, not ext
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_delta_value_serde_round_trip() -> None:
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
original = DeltaValue(
|
||||
delta=[{"type": "human", "content": "hi"}], prev_checkpoint_id="abc-123"
|
||||
)
|
||||
type_tag, blob = serde.dumps_typed(original)
|
||||
assert type_tag == "delta"
|
||||
loaded = serde.loads_typed((type_tag, blob))
|
||||
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
|
||||
|
||||
@@ -12,13 +12,25 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.serde.jsonplus import (
|
||||
JsonPlusSerializer,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
|
||||
|
||||
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:
|
||||
@@ -308,36 +320,3 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
|
||||
assert direct is not None
|
||||
expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict()
|
||||
assert direct.checkpoint["channel_values"]["foo"] == expected
|
||||
|
||||
|
||||
class TestInMemorySaverDeltaChannel:
|
||||
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"
|
||||
version = "00000000000000000000000000000001.0000000000000000"
|
||||
delta = DeltaValue(delta=[{"content": "hi"}], prev_checkpoint_id=None)
|
||||
saver.blobs[(thread_id, ns, channel, version)] = serde.dumps_typed(delta)
|
||||
|
||||
cp = empty_checkpoint()
|
||||
cp["id"] = "cp1"
|
||||
cp["channel_versions"][channel] = version
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp), serde.dumps_typed({}), None)
|
||||
}
|
||||
|
||||
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_blob_missing(self) -> None:
|
||||
"""get_channel_blob returns NotImplemented when checkpoint or channel not found."""
|
||||
saver = InMemorySaver()
|
||||
assert (
|
||||
saver.get_channel_blob("t1", "", "no-such-cp", "messages") is NotImplemented
|
||||
)
|
||||
|
||||
@@ -5,5 +5,5 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.0.1"
|
||||
"langchain-openai==1.1.14"
|
||||
]
|
||||
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.0.0a2",
|
||||
"langchain-openai==1.1.14",
|
||||
"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.0.0a2",
|
||||
"langchain-openai==1.1.14",
|
||||
"langgraph==1.1.2",
|
||||
"langchain_community>=0.3.0",
|
||||
]
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.4.22"
|
||||
__version__ = "0.4.23"
|
||||
|
||||
@@ -23,7 +23,7 @@ dependencies = [
|
||||
path = "langgraph_cli/__init__.py"
|
||||
[project.optional-dependencies]
|
||||
inmem = [
|
||||
"langgraph-api>=0.5.35,<0.8.0 ; python_version >= '3.11'",
|
||||
"langgraph-api>=0.5.35,<0.9.0 ; python_version >= '3.11'",
|
||||
"langgraph-runtime-inmem>=0.7 ; python_version >= '3.11'",
|
||||
]
|
||||
|
||||
|
||||
Generated
+465
-383
File diff suppressed because it is too large
Load Diff
@@ -245,15 +245,6 @@ 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,
|
||||
*,
|
||||
@@ -321,15 +312,6 @@ 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,
|
||||
*,
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from langgraph.channels.any_value import AnyValue
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
|
||||
from langgraph.channels.named_barrier_value import (
|
||||
@@ -21,7 +20,6 @@ __all__ = (
|
||||
"UntrackedValue",
|
||||
"EphemeralValue",
|
||||
"BinaryOperatorAggregate",
|
||||
"DeltaChannel",
|
||||
"NamedBarrierValue",
|
||||
"NamedBarrierValueAfterFinish",
|
||||
# topics
|
||||
|
||||
@@ -119,12 +119,3 @@ 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
|
||||
|
||||
@@ -1,206 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import collections.abc
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel, Value
|
||||
from langgraph.channels.binop import _get_overwrite, _strip_extras
|
||||
from langgraph.errors import EmptyChannelError
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
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).
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
Usage::
|
||||
|
||||
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_every",
|
||||
"_pending",
|
||||
"_base_version",
|
||||
"_last_checkpoint_id",
|
||||
"_overwritten",
|
||||
"_steps_since_snapshot",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[list[Value], Any], list[Value]],
|
||||
typ: type = list,
|
||||
*,
|
||||
snapshot_every: int | None = None,
|
||||
) -> None:
|
||||
typ = _strip_extras(typ)
|
||||
if typ in (
|
||||
collections.abc.Sequence,
|
||||
collections.abc.MutableSequence,
|
||||
):
|
||||
typ = list
|
||||
super().__init__(typ)
|
||||
self.operator = operator
|
||||
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_every != other.snapshot_every:
|
||||
return False
|
||||
if (
|
||||
self.operator.__name__ != "<lambda>"
|
||||
and other.operator.__name__ != "<lambda>"
|
||||
):
|
||||
return self.operator is other.operator
|
||||
return True
|
||||
|
||||
@property
|
||||
def ValueType(self) -> Any:
|
||||
return list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ | list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
|
||||
new.key = self.key
|
||||
new.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 from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
new.value = []
|
||||
elif isinstance(checkpoint, DeltaChainValue):
|
||||
accumulated: list[Value] = list(checkpoint.base) if checkpoint.base else []
|
||||
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 — the pregel layer assembles DeltaValues
|
||||
# into DeltaChainValue before calling from_checkpoint.
|
||||
raise AssertionError(
|
||||
"DeltaChannel.from_checkpoint received a raw DeltaValue. "
|
||||
"This is a bug in the pregel layer — chain assembly should have "
|
||||
"occurred before from_checkpoint was called."
|
||||
)
|
||||
else:
|
||||
# Backwards compat: plain list from old BinaryOperatorAggregate checkpoint.
|
||||
new.value = list(checkpoint)
|
||||
new._pending = []
|
||||
new._base_version = None # set by the subsequent after_checkpoint() call
|
||||
new._overwritten = False
|
||||
return new
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
if not values:
|
||||
return False
|
||||
seen_overwrite = False
|
||||
for value in values:
|
||||
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 = (
|
||||
list(overwrite_value) if overwrite_value is not None else []
|
||||
)
|
||||
self._pending = list(self.value)
|
||||
self._overwritten = True
|
||||
seen_overwrite = True
|
||||
elif not seen_overwrite:
|
||||
self.value = self.operator(self.value, value)
|
||||
self._pending.append(value)
|
||||
return True
|
||||
|
||||
def get(self) -> list[Value]:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
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 list(self.value)
|
||||
return DeltaValue(
|
||||
delta=self._pending[:],
|
||||
prev_checkpoint_id=None if self._overwritten else self._last_checkpoint_id,
|
||||
)
|
||||
|
||||
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
|
||||
@@ -184,39 +184,72 @@ def add_messages(
|
||||
```
|
||||
|
||||
"""
|
||||
remove_all_idx = None
|
||||
# coerce to list
|
||||
if not isinstance(left, list):
|
||||
left = [left] # type: ignore[assignment]
|
||||
if not isinstance(right, list):
|
||||
right = [right] # type: ignore[assignment]
|
||||
# coerce to message
|
||||
left = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(left)
|
||||
]
|
||||
right = [
|
||||
|
||||
# Optimization 1: skip conversion + ID assignment on left when it already
|
||||
# contains fully-resolved BaseMessage objects (the common case after the
|
||||
# first call, since add_messages always returns list[BaseMessage] with IDs).
|
||||
left_msgs: list[BaseMessage]
|
||||
left_seq = cast(list, left)
|
||||
if (
|
||||
left_seq
|
||||
and isinstance(left_seq[0], BaseMessage)
|
||||
and left_seq[0].id is not None
|
||||
and not isinstance(left_seq[0], BaseMessageChunk)
|
||||
):
|
||||
left_msgs = left_seq
|
||||
else:
|
||||
left_msgs = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(left)
|
||||
]
|
||||
for m in left_msgs:
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
|
||||
# always normalise right — it's fresh external input
|
||||
right_msgs: list[BaseMessage] = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(right)
|
||||
]
|
||||
# assign missing ids
|
||||
for m in left:
|
||||
remove_all_idx = None
|
||||
has_remove = False
|
||||
for idx, m in enumerate(right_msgs):
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
for idx, m in enumerate(right):
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:
|
||||
remove_all_idx = idx
|
||||
if isinstance(m, RemoveMessage):
|
||||
has_remove = True
|
||||
if m.id == REMOVE_ALL_MESSAGES:
|
||||
remove_all_idx = idx
|
||||
|
||||
if remove_all_idx is not None:
|
||||
return right[remove_all_idx + 1 :]
|
||||
return right_msgs[remove_all_idx + 1 :]
|
||||
|
||||
# merge
|
||||
merged = left.copy()
|
||||
# Optimization 2: pure-append fast path — no removals, no ID overlaps with
|
||||
# left, and no duplicate IDs within right (all imply a dedup/update is needed).
|
||||
if not has_remove:
|
||||
left_ids = {m.id for m in left_msgs}
|
||||
right_id_set = {m.id for m in right_msgs}
|
||||
if len(right_id_set) == len(right_msgs) and not (right_id_set & left_ids):
|
||||
result = left_msgs + right_msgs
|
||||
if format == "langchain-openai":
|
||||
return _format_messages(result)
|
||||
elif format:
|
||||
msg = (
|
||||
f"Unrecognized {format=}. Expected one of 'langchain-openai', None."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
return result
|
||||
|
||||
# slow path: updates or removals present — full indexed merge
|
||||
merged = left_msgs.copy()
|
||||
merged_by_id = {m.id: i for i, m in enumerate(merged)}
|
||||
ids_to_remove = set()
|
||||
for m in right:
|
||||
for m in right_msgs:
|
||||
if (existing_idx := merged_by_id.get(m.id)) is not None:
|
||||
if isinstance(m, RemoveMessage):
|
||||
ids_to_remove.add(m.id)
|
||||
@@ -228,7 +261,6 @@ def add_messages(
|
||||
raise ValueError(
|
||||
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
|
||||
)
|
||||
|
||||
merged_by_id[m.id] = len(merged)
|
||||
merged.append(m)
|
||||
merged = [m for m in merged if m.id not in ids_to_remove]
|
||||
@@ -238,8 +270,6 @@ def add_messages(
|
||||
elif format:
|
||||
msg = f"Unrecognized {format=}. Expected one of 'langchain-openai', None."
|
||||
raise ValueError(msg)
|
||||
else:
|
||||
pass
|
||||
|
||||
return merged
|
||||
|
||||
|
||||
@@ -1,17 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
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 (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
)
|
||||
from langgraph.checkpoint.base import Checkpoint
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
@@ -20,171 +12,6 @@ from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
LATEST_VERSION = 4
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MISSING_SENTINEL = object()
|
||||
|
||||
|
||||
def _assemble_delta_channels(
|
||||
checkpoint: Checkpoint,
|
||||
config: RunnableConfig,
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
) -> dict[str, Any]:
|
||||
"""Resolve any DeltaValue entries in checkpoint channel_values to DeltaChainValue.
|
||||
|
||||
Returns a dict of only the channels that needed assembly (others are untouched).
|
||||
Tries get_channel_blob fast-path first; falls back to get_tuple traversal.
|
||||
"""
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
current_checkpoint_id = checkpoint.get("id")
|
||||
assembled: dict[str, Any] = {}
|
||||
|
||||
for channel, value in checkpoint["channel_values"].items():
|
||||
if not isinstance(value, DeltaValue):
|
||||
continue
|
||||
|
||||
chain_deltas: list[list[Any]] = []
|
||||
base: list[Any] | None = None
|
||||
cursor: DeltaValue = value
|
||||
# Pre-seed with current checkpoint ID to guard against self-referential chains.
|
||||
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 # chain root
|
||||
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)
|
||||
|
||||
# Fast path: saver has a dedicated blob store.
|
||||
blob = checkpointer.get_channel_blob(
|
||||
thread_id, checkpoint_ns, prev_id, channel
|
||||
)
|
||||
if blob is not NotImplemented:
|
||||
if isinstance(blob, DeltaValue):
|
||||
cursor = blob
|
||||
continue
|
||||
else:
|
||||
base = blob # plain list = snapshot root
|
||||
break
|
||||
|
||||
# Fallback: load the full checkpoint and extract channel value.
|
||||
parent_config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": prev_id,
|
||||
}
|
||||
}
|
||||
parent_tuple = checkpointer.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
|
||||
elif isinstance(prev_val, DeltaValue):
|
||||
cursor = prev_val
|
||||
else:
|
||||
base = prev_val
|
||||
break
|
||||
|
||||
chain_deltas.reverse()
|
||||
assembled[channel] = DeltaChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
return assembled
|
||||
|
||||
|
||||
async def _aassemble_delta_channels(
|
||||
checkpoint: Checkpoint,
|
||||
config: RunnableConfig,
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
) -> dict[str, Any]:
|
||||
"""Async version of _assemble_delta_channels."""
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
current_checkpoint_id = checkpoint.get("id")
|
||||
assembled: dict[str, Any] = {}
|
||||
|
||||
for channel, value in checkpoint["channel_values"].items():
|
||||
if not isinstance(value, DeltaValue):
|
||||
continue
|
||||
|
||||
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 checkpointer.aget_channel_blob(
|
||||
thread_id, checkpoint_ns, prev_id, channel
|
||||
)
|
||||
if blob is not NotImplemented:
|
||||
if isinstance(blob, DeltaValue):
|
||||
cursor = blob
|
||||
continue
|
||||
else:
|
||||
base = blob
|
||||
break
|
||||
|
||||
parent_config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": prev_id,
|
||||
}
|
||||
}
|
||||
parent_tuple = await checkpointer.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
|
||||
elif isinstance(prev_val, DeltaValue):
|
||||
cursor = prev_val
|
||||
else:
|
||||
base = prev_val
|
||||
break
|
||||
|
||||
chain_deltas.reverse()
|
||||
assembled[channel] = DeltaChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
return assembled
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
@@ -240,12 +67,13 @@ def channels_from_checkpoint(
|
||||
channel_specs[k] = v
|
||||
else:
|
||||
managed_specs[k] = v
|
||||
channels: dict[str, BaseChannel] = {}
|
||||
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
|
||||
return (
|
||||
{
|
||||
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
for k, v in channel_specs.items()
|
||||
},
|
||||
managed_specs,
|
||||
)
|
||||
|
||||
|
||||
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
|
||||
@@ -92,8 +92,6 @@ from langgraph.pregel._algo import (
|
||||
task_path_str,
|
||||
)
|
||||
from langgraph.pregel._checkpoint import (
|
||||
_aassemble_delta_channels,
|
||||
_assemble_delta_channels,
|
||||
channels_from_checkpoint,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
@@ -833,8 +831,18 @@ 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 self.is_replaying:
|
||||
if is_time_traveling:
|
||||
replay_checkpoint_id = self.checkpoint["id"]
|
||||
if (
|
||||
self.checkpoint_metadata.get("source")
|
||||
@@ -883,12 +891,6 @@ class PregelLoop:
|
||||
id=self.checkpoint["id"] if exiting else None,
|
||||
updated_channels=self.updated_channels,
|
||||
)
|
||||
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()
|
||||
@@ -1270,19 +1272,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
else []
|
||||
)
|
||||
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
|
||||
# Assemble any DeltaChannel chains before constructing channel objects.
|
||||
if self.checkpointer is not None:
|
||||
assembled = _assemble_delta_channels(
|
||||
self.checkpoint, self.checkpoint_config, self.checkpointer
|
||||
)
|
||||
if assembled:
|
||||
self.checkpoint = {
|
||||
**self.checkpoint,
|
||||
"channel_values": {
|
||||
**self.checkpoint["channel_values"],
|
||||
**assembled,
|
||||
},
|
||||
}
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs, self.checkpoint
|
||||
)
|
||||
@@ -1487,18 +1476,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
self.submit = await self.stack.enter_async_context(
|
||||
AsyncBackgroundExecutor(self.config)
|
||||
)
|
||||
if self.checkpointer is not None:
|
||||
assembled = await _aassemble_delta_channels(
|
||||
self.checkpoint, self.checkpoint_config, self.checkpointer
|
||||
)
|
||||
if assembled:
|
||||
self.checkpoint = {
|
||||
**self.checkpoint,
|
||||
"channel_values": {
|
||||
**self.checkpoint["channel_values"],
|
||||
**assembled,
|
||||
},
|
||||
}
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs, self.checkpoint
|
||||
)
|
||||
|
||||
@@ -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_END, NS_SEP
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
from langgraph.types import Command
|
||||
@@ -132,23 +132,15 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||
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]
|
||||
ns = tuple(cast(str, metadata["langgraph_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")]:
|
||||
stream_metadata["tags"] = filtered_tags
|
||||
self.metadata[run_id] = (ns, stream_metadata)
|
||||
metadata["tags"] = filtered_tags
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
|
||||
def on_llm_new_token(
|
||||
self,
|
||||
|
||||
@@ -122,8 +122,6 @@ from langgraph.pregel._algo import (
|
||||
)
|
||||
from langgraph.pregel._call import identifier
|
||||
from langgraph.pregel._checkpoint import (
|
||||
_aassemble_delta_channels,
|
||||
_assemble_delta_channels,
|
||||
channels_from_checkpoint,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
@@ -1051,23 +1049,13 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver):
|
||||
assembled = _assemble_delta_channels(
|
||||
checkpoint, saved.config, self.checkpointer
|
||||
)
|
||||
if assembled:
|
||||
checkpoint = {
|
||||
**checkpoint,
|
||||
"channel_values": {**checkpoint["channel_values"], **assembled},
|
||||
}
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saved.checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
saved.checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1180,23 +1168,13 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver):
|
||||
assembled = await _aassemble_delta_channels(
|
||||
checkpoint, saved.config, self.checkpointer
|
||||
)
|
||||
if assembled:
|
||||
checkpoint = {
|
||||
**checkpoint,
|
||||
"channel_values": {**checkpoint["channel_values"], **assembled},
|
||||
}
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saved.checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
saved.checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1542,20 +1520,9 @@ class Pregel(
|
||||
saved = checkpointer.get_tuple(config)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
if saved:
|
||||
assembled = _assemble_delta_channels(
|
||||
base_checkpoint, saved.config, checkpointer
|
||||
)
|
||||
if assembled:
|
||||
base_checkpoint = {
|
||||
**base_checkpoint,
|
||||
"channel_values": {
|
||||
**base_checkpoint["channel_values"],
|
||||
**assembled,
|
||||
},
|
||||
}
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
@@ -1999,20 +1966,9 @@ class Pregel(
|
||||
saved = await checkpointer.aget_tuple(config)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
if saved:
|
||||
assembled = await _aassemble_delta_channels(
|
||||
base_checkpoint, saved.config, checkpointer
|
||||
)
|
||||
if assembled:
|
||||
base_checkpoint = {
|
||||
**base_checkpoint,
|
||||
"channel_values": {
|
||||
**base_checkpoint["channel_values"],
|
||||
**assembled,
|
||||
},
|
||||
}
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.1.7a2"
|
||||
version = "1.1.9"
|
||||
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.0a2",
|
||||
"langchain-core>=1.3.0,<2",
|
||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||
"langgraph-sdk>=0.3.0,<0.4.0",
|
||||
"langgraph-prebuilt>=1.0.9,<1.1.0",
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Benchmark: add_messages fast-path optimizations.
|
||||
|
||||
Both implementations are inlined so the benchmark is self-contained and
|
||||
immune to import-cache or installed-vs-local confusion.
|
||||
|
||||
Run directly:
|
||||
python tests/test_add_messages_benchmark.py
|
||||
|
||||
Or via pytest (correctness only, numbers printed to stdout):
|
||||
pytest tests/test_add_messages_benchmark.py -s -v
|
||||
"""
|
||||
|
||||
import statistics
|
||||
import time
|
||||
import tracemalloc
|
||||
import uuid
|
||||
from typing import cast
|
||||
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
BaseMessageChunk,
|
||||
HumanMessage,
|
||||
RemoveMessage,
|
||||
convert_to_messages,
|
||||
message_chunk_to_message,
|
||||
)
|
||||
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES
|
||||
|
||||
# ── original implementation (pre-optimisation) ────────────────────────────────
|
||||
|
||||
|
||||
def _add_messages_original(left, right):
|
||||
remove_all_idx = None
|
||||
if not isinstance(left, list):
|
||||
left = [left]
|
||||
if not isinstance(right, list):
|
||||
right = [right]
|
||||
left = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(left)
|
||||
]
|
||||
right = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(right)
|
||||
]
|
||||
for m in left:
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
for idx, m in enumerate(right):
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:
|
||||
remove_all_idx = idx
|
||||
if remove_all_idx is not None:
|
||||
return right[remove_all_idx + 1 :]
|
||||
merged = left.copy()
|
||||
merged_by_id = {m.id: i for i, m in enumerate(merged)}
|
||||
ids_to_remove = set()
|
||||
for m in right:
|
||||
if (existing_idx := merged_by_id.get(m.id)) is not None:
|
||||
if isinstance(m, RemoveMessage):
|
||||
ids_to_remove.add(m.id)
|
||||
else:
|
||||
ids_to_remove.discard(m.id)
|
||||
merged[existing_idx] = m
|
||||
else:
|
||||
if isinstance(m, RemoveMessage):
|
||||
raise ValueError(
|
||||
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
|
||||
)
|
||||
merged_by_id[m.id] = len(merged)
|
||||
merged.append(m)
|
||||
return [m for m in merged if m.id not in ids_to_remove]
|
||||
|
||||
|
||||
# ── optimised implementation ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _add_messages_optimized(left, right):
|
||||
if not isinstance(left, list):
|
||||
left = [left]
|
||||
if not isinstance(right, list):
|
||||
right = [right]
|
||||
|
||||
# Optimisation 1: skip conversion + ID assignment on left when it already
|
||||
# contains fully-resolved BaseMessage objects (the common case after the
|
||||
# first call, since add_messages always returns list[BaseMessage] with IDs).
|
||||
if (
|
||||
left
|
||||
and isinstance(left[0], BaseMessage)
|
||||
and not isinstance(left[0], BaseMessageChunk)
|
||||
):
|
||||
left = cast(list[BaseMessage], left)
|
||||
else:
|
||||
left = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(left)
|
||||
]
|
||||
for m in left:
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
|
||||
# always normalise right — it's fresh external input
|
||||
right = [
|
||||
message_chunk_to_message(cast(BaseMessageChunk, m))
|
||||
for m in convert_to_messages(right)
|
||||
]
|
||||
remove_all_idx = None
|
||||
has_remove = False
|
||||
for idx, m in enumerate(right):
|
||||
if m.id is None:
|
||||
m.id = str(uuid.uuid4())
|
||||
if isinstance(m, RemoveMessage):
|
||||
has_remove = True
|
||||
if m.id == REMOVE_ALL_MESSAGES:
|
||||
remove_all_idx = idx
|
||||
|
||||
if remove_all_idx is not None:
|
||||
return right[remove_all_idx + 1 :]
|
||||
|
||||
# Optimisation 2: pure-append fast path — no removals and no ID overlaps.
|
||||
# Builds one set over left instead of copying left + building a full dict.
|
||||
if not has_remove:
|
||||
left_ids = {m.id for m in left}
|
||||
if not any(m.id in left_ids for m in right):
|
||||
return left + right
|
||||
|
||||
# slow path: updates or removals present — full indexed merge
|
||||
merged = left.copy()
|
||||
merged_by_id = {m.id: i for i, m in enumerate(merged)}
|
||||
ids_to_remove = set()
|
||||
for m in right:
|
||||
if (existing_idx := merged_by_id.get(m.id)) is not None:
|
||||
if isinstance(m, RemoveMessage):
|
||||
ids_to_remove.add(m.id)
|
||||
else:
|
||||
ids_to_remove.discard(m.id)
|
||||
merged[existing_idx] = m
|
||||
else:
|
||||
if isinstance(m, RemoveMessage):
|
||||
raise ValueError(
|
||||
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
|
||||
)
|
||||
merged_by_id[m.id] = len(merged)
|
||||
merged.append(m)
|
||||
return [m for m in merged if m.id not in ids_to_remove]
|
||||
|
||||
|
||||
# ── helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_messages(n: int) -> list[BaseMessage]:
|
||||
return [
|
||||
(HumanMessage if i % 2 == 0 else AIMessage)(
|
||||
content=f"message {i}", id=str(uuid.uuid4())
|
||||
)
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
|
||||
def _bench_time(fn, left, right, *, iters: int = 2_000) -> float:
|
||||
"""Return median latency in microseconds."""
|
||||
for _ in range(100):
|
||||
fn(list(left), list(right))
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
left_copy, right_copy = list(left), list(right)
|
||||
t0 = time.perf_counter()
|
||||
fn(left_copy, right_copy)
|
||||
times.append(time.perf_counter() - t0)
|
||||
return statistics.median(times) * 1e6
|
||||
|
||||
|
||||
def _bench_memory(fn, left, right) -> int:
|
||||
"""Return peak memory allocated during a single call (bytes)."""
|
||||
# one warm-up so any lazy init is excluded
|
||||
fn(list(left), list(right))
|
||||
left_copy, right_copy = list(left), list(right)
|
||||
tracemalloc.start()
|
||||
tracemalloc.clear_traces()
|
||||
fn(left_copy, right_copy)
|
||||
_, peak = tracemalloc.get_traced_memory()
|
||||
tracemalloc.stop()
|
||||
return peak
|
||||
|
||||
|
||||
# ── scenarios ─────────────────────────────────────────────────────────────────
|
||||
|
||||
SCENARIOS = [
|
||||
("pure append 1 → 1 msg", 1, 1, "append"),
|
||||
("pure append 10 → 1 msg", 10, 1, "append"),
|
||||
("pure append 100 → 1 msg", 100, 1, "append"),
|
||||
("pure append 1000 → 1 msg", 1000, 1, "append"),
|
||||
("pure append 1000 → 5 msgs", 1000, 5, "append"),
|
||||
("update existing 100 → 1 msg", 100, 1, "update"),
|
||||
("remove message 100 → 1 msg", 100, 1, "remove"),
|
||||
]
|
||||
|
||||
|
||||
def _make_inputs(n_left, n_right, mode):
|
||||
left = _make_messages(n_left)
|
||||
right = _make_messages(n_right)
|
||||
if mode == "update":
|
||||
right[0] = AIMessage(content="updated", id=left[0].id)
|
||||
elif mode == "remove":
|
||||
right = [RemoveMessage(id=left[0].id)]
|
||||
return left, right
|
||||
|
||||
|
||||
# ── main output ───────────────────────────────────────────────────────────────
|
||||
|
||||
COL = 36
|
||||
|
||||
|
||||
def run_benchmarks() -> None:
|
||||
print()
|
||||
print("=" * 88)
|
||||
print("add_messages benchmark — time (µs, median of 2 000 iterations)")
|
||||
print("=" * 88)
|
||||
print(f"{'Scenario':<{COL}} {'Original':>10} {'Optimized':>11} {'Speedup':>8}")
|
||||
print("-" * 88)
|
||||
|
||||
for label, n_left, n_right, mode in SCENARIOS:
|
||||
left, right = _make_inputs(n_left, n_right, mode)
|
||||
t_orig = _bench_time(_add_messages_original, left, right)
|
||||
t_opt = _bench_time(_add_messages_optimized, left, right)
|
||||
print(f"{label:<{COL}} {t_orig:>10.2f} {t_opt:>11.2f} {t_orig / t_opt:>7.2f}x")
|
||||
|
||||
print()
|
||||
print("=" * 88)
|
||||
print("add_messages benchmark — peak memory allocated per call (bytes)")
|
||||
print("=" * 88)
|
||||
print(f"{'Scenario':<{COL}} {'Original':>10} {'Optimized':>11} {'Reduction':>10}")
|
||||
print("-" * 88)
|
||||
|
||||
for label, n_left, n_right, mode in SCENARIOS:
|
||||
left, right = _make_inputs(n_left, n_right, mode)
|
||||
m_orig = _bench_memory(_add_messages_original, left, right)
|
||||
m_opt = _bench_memory(_add_messages_optimized, left, right)
|
||||
reduction = (1 - m_opt / m_orig) * 100 if m_orig else 0.0
|
||||
print(f"{label:<{COL}} {m_orig:>10,} {m_opt:>11,} {reduction:>9.1f}%")
|
||||
|
||||
print()
|
||||
print("=" * 88)
|
||||
print("Simulated long thread — 200 steps × 2 msgs appended per step")
|
||||
print("=" * 88)
|
||||
for name, fn in [
|
||||
("original", _add_messages_original),
|
||||
("optimized", _add_messages_optimized),
|
||||
]:
|
||||
state: list = []
|
||||
t0 = time.perf_counter()
|
||||
for step in range(200):
|
||||
new_msgs = [
|
||||
HumanMessage(content=f"step {step} human", id=str(uuid.uuid4())),
|
||||
AIMessage(content=f"step {step} ai", id=str(uuid.uuid4())),
|
||||
]
|
||||
state = fn(state, new_msgs)
|
||||
elapsed = (time.perf_counter() - t0) * 1_000
|
||||
print(f" {name:<12} {elapsed:.2f} ms ({len(state)} messages)")
|
||||
print()
|
||||
|
||||
|
||||
# ── pytest entry-points ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_add_messages_correctness():
|
||||
"""Optimised implementation must match original output for every scenario."""
|
||||
for label, n_left, n_right, mode in SCENARIOS:
|
||||
left, right = _make_inputs(n_left, n_right, mode)
|
||||
expected = _add_messages_original(list(left), list(right))
|
||||
actual = _add_messages_optimized(list(left), list(right))
|
||||
assert len(actual) == len(expected), f"[{label}] length mismatch"
|
||||
for a, b in zip(actual, expected):
|
||||
assert type(a) is type(b), f"[{label}] type mismatch"
|
||||
assert a.id == b.id, f"[{label}] id mismatch"
|
||||
assert a.content == b.content, f"[{label}] content mismatch"
|
||||
|
||||
|
||||
def test_add_messages_benchmark(capsys):
|
||||
run_benchmarks()
|
||||
out = capsys.readouterr().out
|
||||
assert "Speedup" in out
|
||||
assert "Optimized" in out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_benchmarks()
|
||||
@@ -117,408 +117,3 @@ def test_untracked_value() -> None:
|
||||
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
|
||||
with pytest.raises(EmptyChannelError):
|
||||
new_channel.get()
|
||||
|
||||
|
||||
def test_delta_channel_basic_two_steps() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
# Step 1: one message added
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
assert len(d1.delta) == 1
|
||||
assert d1.prev_checkpoint_id is None # first ever step
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
# Step 2: another message
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert d2.prev_checkpoint_id == "cid1"
|
||||
assert len(d2.delta) == 1
|
||||
ch.after_checkpoint("v2")
|
||||
|
||||
# Full accumulated value is preserved in memory
|
||||
assert len(ch.get()) == 2
|
||||
assert ch.get()[0].content == "hi"
|
||||
assert ch.get()[1].content == "hello"
|
||||
|
||||
|
||||
def test_delta_channel_after_checkpoint_no_op_when_unchanged() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
ch.after_checkpoint("v1")
|
||||
|
||||
# Same version: no-op
|
||||
ch.after_checkpoint("v1")
|
||||
assert ch._base_version == "v1"
|
||||
assert ch._pending == []
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_chain() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
chain = DeltaChainValue(
|
||||
base=None,
|
||||
deltas=[
|
||||
[HumanMessage(content="hi", id="h1")],
|
||||
[AIMessage(content="hello", id="a1")],
|
||||
[HumanMessage(content="bye", id="h2")],
|
||||
],
|
||||
)
|
||||
ch = spec.from_checkpoint(chain)
|
||||
msgs = ch.get()
|
||||
assert len(msgs) == 3
|
||||
assert msgs[0].content == "hi"
|
||||
assert msgs[1].content == "hello"
|
||||
assert msgs[2].content == "bye"
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# Old BinaryOperatorAggregate checkpoint: plain list
|
||||
spec = DeltaChannel(add_messages)
|
||||
old_value = [HumanMessage(content="old", id="h1")]
|
||||
ch = spec.from_checkpoint(old_value)
|
||||
assert ch.get() == old_value
|
||||
|
||||
|
||||
def test_delta_channel_overwrite_resets_chain() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
ch.update([HumanMessage(content="old", id="h1")])
|
||||
ch.after_checkpoint("v1")
|
||||
|
||||
# Overwrite should create a root blob (prev_checkpoint_id=None)
|
||||
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
|
||||
d = ch.checkpoint()
|
||||
assert isinstance(d, DeltaValue)
|
||||
assert d.prev_checkpoint_id is None # chain root
|
||||
assert len(d.delta) == 1
|
||||
assert d.delta[0].content == "new"
|
||||
|
||||
|
||||
def test_delta_channel_assembly_fallback_via_get_tuple() -> None:
|
||||
"""Assembly falls back to get_tuple for savers without get_channel_blob."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
CheckpointTuple,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.pregel._checkpoint import _assemble_delta_channels
|
||||
|
||||
msg1 = {"type": "human", "content": "hello"}
|
||||
msg2 = {"type": "ai", "content": "world"}
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp1["channel_values"]["messages"] = [msg1]
|
||||
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "cp2"
|
||||
cp2["channel_values"]["messages"] = DeltaValue(
|
||||
delta=[msg2], prev_checkpoint_id="cp1"
|
||||
)
|
||||
|
||||
saver = MagicMock()
|
||||
saver.get_channel_blob.return_value = NotImplemented
|
||||
saver.get_tuple.return_value = CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "t1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "cp1",
|
||||
}
|
||||
},
|
||||
checkpoint=cp1,
|
||||
metadata={},
|
||||
parent_config=None,
|
||||
pending_writes=[],
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
|
||||
assembled = _assemble_delta_channels(cp2, config, saver)
|
||||
|
||||
assert "messages" in assembled
|
||||
chain = assembled["messages"]
|
||||
assert isinstance(chain, DeltaChainValue)
|
||||
assert chain.base == [msg1]
|
||||
assert chain.deltas == [[msg2]]
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(chain)
|
||||
result = ch.get()
|
||||
assert len(result) == 2
|
||||
assert isinstance(result[0], HumanMessage) and result[0].content == "hello"
|
||||
assert isinstance(result[1], AIMessage) and result[1].content == "world"
|
||||
|
||||
|
||||
def test_delta_channel_remove_message_delta_and_replay() -> None:
|
||||
"""RemoveMessage stored in a delta must round-trip correctly through the chain."""
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
# Step 1: add two messages
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
assert ch.get() == [
|
||||
HumanMessage(content="hi", id="h1"),
|
||||
AIMessage(content="hello", id="a1"),
|
||||
]
|
||||
|
||||
# Step 2: remove the AI message
|
||||
ch.update([RemoveMessage(id="a1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert isinstance(d2, DeltaValue)
|
||||
assert d2.prev_checkpoint_id == "cid1"
|
||||
assert any(isinstance(w, RemoveMessage) for w in d2.delta)
|
||||
ch.after_checkpoint("v2", checkpoint_id="cid2")
|
||||
assert ch.get() == [HumanMessage(content="hi", id="h1")]
|
||||
|
||||
# Replay the full chain from scratch — must reproduce the post-remove state
|
||||
chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta])
|
||||
ch2 = spec.from_checkpoint(chain)
|
||||
assert ch2.get() == [HumanMessage(content="hi", id="h1")]
|
||||
|
||||
|
||||
def test_delta_channel_update_by_id_delta_and_replay() -> None:
|
||||
"""Updating a message by ID stored in a delta must round-trip correctly."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
# Step 1: add a message
|
||||
ch.update([HumanMessage(content="original", id="h1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
# Step 2: update the same message by ID
|
||||
ch.update([HumanMessage(content="updated", id="h1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert isinstance(d2, DeltaValue)
|
||||
assert d2.prev_checkpoint_id == "cid1"
|
||||
ch.after_checkpoint("v2", checkpoint_id="cid2")
|
||||
assert ch.get() == [HumanMessage(content="updated", id="h1")]
|
||||
|
||||
# Replay the full chain — must produce the updated message, not the original
|
||||
chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta])
|
||||
ch2 = spec.from_checkpoint(chain)
|
||||
assert len(ch2.get()) == 1
|
||||
assert ch2.get()[0].content == "updated"
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_every_emits_plain_list() -> None:
|
||||
"""snapshot_every=N causes a plain-list snapshot after N steps; next deltas chain to it."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
SNAP = 3
|
||||
spec = DeltaChannel(add_messages, snapshot_every=SNAP)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
# First after_checkpoint anchors _base_version without counting a step.
|
||||
ch.after_checkpoint("v0", checkpoint_id="cid0")
|
||||
|
||||
# Steps 1..SNAP: each should stay as DeltaValue; counter increments each step.
|
||||
for i in range(1, SNAP + 1):
|
||||
ch.update([HumanMessage(content=f"m{i}", id=f"h{i}")])
|
||||
ckpt = ch.checkpoint()
|
||||
assert isinstance(ckpt, DeltaValue), f"expected DeltaValue at step {i}"
|
||||
ch.after_checkpoint(f"v{i}", checkpoint_id=f"cid{i}")
|
||||
|
||||
# Step SNAP+1: _steps_since_snapshot == SNAP → snapshot fires
|
||||
ch.update([HumanMessage(content="snap", id="hsnap")])
|
||||
snap = ch.checkpoint()
|
||||
assert isinstance(snap, list), "expected plain-list snapshot at snapshot_every step"
|
||||
assert len(snap) == SNAP + 1
|
||||
|
||||
# After snapshot, counter resets — next step is DeltaValue again
|
||||
ch.after_checkpoint("vsnap", checkpoint_id="cidsnap")
|
||||
ch.update([HumanMessage(content="post", id="hpost")])
|
||||
post = ch.checkpoint()
|
||||
assert isinstance(post, DeltaValue)
|
||||
assert post.prev_checkpoint_id == "cidsnap"
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_every_end_to_end() -> None:
|
||||
"""Graph with snapshot_every: get_state returns correct accumulated value after snapshot."""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)]
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
counter["n"] += 1
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
|
||||
]
|
||||
}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "snap-test"}}
|
||||
|
||||
# Run 5 turns — snapshot fires after 2 steps, then again after 2 more
|
||||
for i in range(5):
|
||||
graph.invoke({"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# 5 human + 5 AI = 10 total
|
||||
assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}"
|
||||
|
||||
|
||||
def test_delta_channel_assembly_fast_path_returns_delta_value() -> None:
|
||||
"""get_channel_blob returning a DeltaValue continues chain traversal (fast-path)."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.pregel._checkpoint import _assemble_delta_channels
|
||||
|
||||
msg1 = {"type": "human", "content": "one"}
|
||||
msg2 = {"type": "ai", "content": "two"}
|
||||
msg3 = {"type": "human", "content": "three"}
|
||||
|
||||
# cp3 → cp2 (DeltaValue) → cp1 (base list)
|
||||
dv_cp2 = DeltaValue(delta=[msg2], prev_checkpoint_id="cp1")
|
||||
cp3 = empty_checkpoint()
|
||||
cp3["id"] = "cp3"
|
||||
cp3["channel_values"]["messages"] = DeltaValue(
|
||||
delta=[msg3], prev_checkpoint_id="cp2"
|
||||
)
|
||||
|
||||
saver = MagicMock()
|
||||
|
||||
def _get_blob(thread_id, ns, checkpoint_id, channel):
|
||||
if checkpoint_id == "cp2":
|
||||
return dv_cp2 # DeltaValue — chain continues
|
||||
if checkpoint_id == "cp1":
|
||||
return [msg1] # plain list — chain root
|
||||
return NotImplemented
|
||||
|
||||
saver.get_channel_blob.side_effect = _get_blob
|
||||
|
||||
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
|
||||
assembled = _assemble_delta_channels(cp3, config, saver)
|
||||
|
||||
chain = assembled["messages"]
|
||||
assert isinstance(chain, DeltaChainValue)
|
||||
assert chain.base == [msg1]
|
||||
assert chain.deltas == [[msg2], [msg3]]
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(chain)
|
||||
# add_messages converts dicts to message objects; check by type and content
|
||||
|
||||
result = ch.get()
|
||||
assert len(result) == 3
|
||||
assert result[0].content == "one"
|
||||
assert result[1].content == "two"
|
||||
assert result[2].content == "three"
|
||||
|
||||
|
||||
def test_delta_channel_assembly_broken_chain_logs_warning() -> None:
|
||||
"""If a prev_checkpoint_id points to a missing checkpoint, log a warning and use partial chain."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langgraph.checkpoint.base import DeltaValue, empty_checkpoint
|
||||
|
||||
from langgraph.pregel._checkpoint import _assemble_delta_channels
|
||||
|
||||
cp = empty_checkpoint()
|
||||
cp["id"] = "cp2"
|
||||
cp["channel_values"]["messages"] = DeltaValue(
|
||||
delta=["msg2"], prev_checkpoint_id="cp-missing"
|
||||
)
|
||||
|
||||
saver = MagicMock()
|
||||
saver.get_channel_blob.return_value = NotImplemented
|
||||
saver.get_tuple.return_value = None # checkpoint not found
|
||||
|
||||
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
|
||||
|
||||
assembled = _assemble_delta_channels(cp, config, saver)
|
||||
|
||||
# Should still assemble — with partial chain (just the current delta, base=None)
|
||||
assert "messages" in assembled
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
|
||||
chain = assembled["messages"]
|
||||
assert isinstance(chain, DeltaChainValue)
|
||||
assert chain.base is None
|
||||
assert chain.deltas == [["msg2"]]
|
||||
|
||||
@@ -1,356 +0,0 @@
|
||||
"""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
|
||||
|
||||
Simulates realistic multi-turn conversations with paragraph-length messages
|
||||
(~100 tokens each) scaling up to 1M-token-equivalent histories.
|
||||
|
||||
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 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
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
|
||||
_SQLITE_AVAILABLE = True
|
||||
except ImportError:
|
||||
_SQLITE_AVAILABLE = False
|
||||
|
||||
SNAPSHOT_EVERY = 50
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Realistic message payload (~100 tokens / ~400 chars each)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_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."
|
||||
)
|
||||
|
||||
_AI_TEMPLATE = (
|
||||
"Great question about {topic}. The key insight here is that {concern} introduces "
|
||||
"a subtle ordering dependency that most teams overlook until they hit it in production. "
|
||||
"For your {component} layer specifically, I'd recommend starting with a careful audit "
|
||||
"of the interface boundaries before making any structural changes. This will give you "
|
||||
"a clear picture of the blast radius and let you sequence the migration safely."
|
||||
)
|
||||
|
||||
_TOPICS = [
|
||||
"distributed tracing",
|
||||
"eventual consistency",
|
||||
"schema migration",
|
||||
"backpressure handling",
|
||||
"idempotency guarantees",
|
||||
"cache invalidation",
|
||||
"connection pooling",
|
||||
"rate limiting",
|
||||
"circuit breaking",
|
||||
"observability pipelines",
|
||||
]
|
||||
|
||||
_CONCERNS = [
|
||||
"concurrency model",
|
||||
"retry semantics",
|
||||
"state management",
|
||||
"error propagation",
|
||||
"latency budget",
|
||||
]
|
||||
|
||||
_COMPONENTS = [
|
||||
"persistence",
|
||||
"routing",
|
||||
"ingestion",
|
||||
"aggregation",
|
||||
"serialization",
|
||||
]
|
||||
|
||||
|
||||
def _human_content(i: int) -> str:
|
||||
return _HUMAN_TEMPLATE.format(
|
||||
topic=_TOPICS[i % len(_TOPICS)],
|
||||
concern=_CONCERNS[i % len(_CONCERNS)],
|
||||
component=_COMPONENTS[i % len(_COMPONENTS)],
|
||||
)
|
||||
|
||||
|
||||
def _ai_content(i: int) -> str:
|
||||
return _AI_TEMPLATE.format(
|
||||
topic=_TOPICS[i % len(_TOPICS)],
|
||||
concern=_CONCERNS[i % len(_CONCERNS)],
|
||||
component=_COMPONENTS[i % len(_COMPONENTS)],
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# State definitions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BinaryState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
|
||||
|
||||
class DeltaState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
|
||||
class DeltaSnapshotState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=SNAPSHOT_EVERY)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_graph(state_cls: type, checkpointer: Any = None) -> Any:
|
||||
def human_node(state: Any) -> dict:
|
||||
return {}
|
||||
|
||||
def ai_node(state: Any) -> dict:
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=_ai_content(i), id=f"a{i}")]}
|
||||
|
||||
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=checkpointer or MemorySaver())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Measurement helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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 _run_turns(
|
||||
n_turns: int,
|
||||
state_cls: type,
|
||||
checkpointer: Any = None,
|
||||
) -> tuple[float, float, int]:
|
||||
"""Run n_turns conversation turns.
|
||||
|
||||
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
|
||||
blob_bytes is -1 for savers without in-memory blob stores (e.g. SQLite).
|
||||
Read latency is measured as the time to invoke the graph with no new
|
||||
messages after the full history is built — this forces state rehydration.
|
||||
"""
|
||||
graph = _make_graph(state_cls, checkpointer)
|
||||
config = {"configurable": {"thread_id": "bench"}}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
for i in range(n_turns):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
|
||||
config,
|
||||
)
|
||||
write_elapsed = time.perf_counter() - t0
|
||||
|
||||
# Measure read/rehydration: get_state forces the channel to rebuild
|
||||
t1 = time.perf_counter()
|
||||
for _ in range(5):
|
||||
graph.get_state(config)
|
||||
read_elapsed = (time.perf_counter() - t1) / 5
|
||||
|
||||
if isinstance(graph.checkpointer, MemorySaver):
|
||||
blob_bytes = _total_blob_bytes(graph.checkpointer)
|
||||
else:
|
||||
blob_bytes = -1
|
||||
return write_elapsed, read_elapsed, blob_bytes
|
||||
|
||||
|
||||
def _fmt_bytes(n: int) -> str:
|
||||
if n >= 1_000_000:
|
||||
return f"{n / 1_000_000:.1f} MB"
|
||||
if n >= 1_000:
|
||||
return f"{n / 1_000:.1f} KB"
|
||||
return f"{n} B"
|
||||
|
||||
|
||||
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"
|
||||
if tokens >= 1_000:
|
||||
return f"~{tokens / 1_000:.0f}K tok"
|
||||
return f"~{tokens} tok"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# 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 _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
|
||||
|
||||
factories.append(("SQLite", tempfile.NamedTemporaryFile(suffix=".db")))
|
||||
return factories
|
||||
|
||||
|
||||
def run_benchmark() -> None:
|
||||
print()
|
||||
print(
|
||||
"DeltaChannel vs add_messages (BinaryOperatorAggregate) — checkpoint storage & latency"
|
||||
)
|
||||
print("Simulating realistic multi-turn conversations up to ~1M-token histories")
|
||||
print("(5,000 turns × ~200 tokens/turn ≈ 1M tokens — Claude's full context window)")
|
||||
print()
|
||||
|
||||
checkpointers: list[tuple[str, Any]] = [("InMemory (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_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
import contextlib
|
||||
import tempfile
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
yield None
|
||||
else:
|
||||
with tempfile.NamedTemporaryFile(suffix=".db") as f:
|
||||
with SqliteSaver.from_conn_string(f.name) as saver:
|
||||
yield saver
|
||||
|
||||
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)
|
||||
with _make_saver() as saver:
|
||||
s_wt, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
|
||||
|
||||
# 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:
|
||||
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"{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("=" * 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" Storage: {_fmt_bytes(b_bytes)} (add_messages) → {_fmt_bytes(s_bytes)} (DeltaChannel+snapshot) — {ratio:.0f}x reduction"
|
||||
)
|
||||
print(
|
||||
f" Read latency: {b_rt * 1000:.1f}ms (add_messages) vs {s_rt * 1000:.1f}ms (DeltaChannel+snapshot)"
|
||||
)
|
||||
print()
|
||||
print("Legend:")
|
||||
print(
|
||||
" 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()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_channel_benchmark(capsys: Any) -> None:
|
||||
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
|
||||
with capsys.disabled():
|
||||
run_benchmark()
|
||||
|
||||
# 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}"
|
||||
)
|
||||
assert s_bytes < b_bytes, (
|
||||
f"DeltaChannel+snapshot should use less storage at {turns} turns, "
|
||||
f"got snapshot={s_bytes} binary={b_bytes}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Script entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_benchmark()
|
||||
sys.exit(0)
|
||||
@@ -275,3 +275,70 @@ 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 == []
|
||||
|
||||
@@ -5,6 +5,7 @@ import langchain_core
|
||||
import pytest
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
AIMessageChunk,
|
||||
AnyMessage,
|
||||
HumanMessage,
|
||||
RemoveMessage,
|
||||
@@ -338,6 +339,123 @@ def test_remove_all_messages():
|
||||
]
|
||||
|
||||
|
||||
def test_fast_path_preserves_format_openai():
|
||||
"""Pure-append fast path must still apply the `langchain-openai` formatter."""
|
||||
left = [HumanMessage(content="prior", id="1")]
|
||||
right = [
|
||||
AIMessage(
|
||||
content=[
|
||||
{
|
||||
"type": "tool_use",
|
||||
"name": "foo",
|
||||
"input": {"bar": "baz"},
|
||||
"id": "t1",
|
||||
}
|
||||
],
|
||||
id="2",
|
||||
)
|
||||
]
|
||||
result = add_messages(left, right, format="langchain-openai")
|
||||
assert isinstance(result[0], HumanMessage)
|
||||
assert result[0].content == "prior"
|
||||
assert isinstance(result[1], AIMessage)
|
||||
# formatter collapses the tool_use content block into `tool_calls`
|
||||
assert result[1].content == ""
|
||||
assert len(result[1].tool_calls) == 1
|
||||
assert result[1].tool_calls[0]["name"] == "foo"
|
||||
assert result[1].tool_calls[0]["args"] == {"bar": "baz"}
|
||||
assert result[1].tool_calls[0]["id"] == "t1"
|
||||
|
||||
|
||||
def test_fast_path_rejects_invalid_format():
|
||||
"""Pure-append fast path must validate the `format` arg like the slow path."""
|
||||
left = [HumanMessage(content="prior", id="1")]
|
||||
right = [AIMessage(content="new", id="2")]
|
||||
with pytest.raises(ValueError, match="Unrecognized format="):
|
||||
add_messages(left, right, format="bogus") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_left_starting_with_chunk_is_normalized():
|
||||
"""Opt-1 guard: a `BaseMessageChunk` at left[0] must trigger full conversion."""
|
||||
chunk = AIMessageChunk(content="chunk", id="c1")
|
||||
result = add_messages([chunk], [HumanMessage(content="h", id="h1")])
|
||||
assert len(result) == 2
|
||||
# chunk must be converted to a non-chunk message
|
||||
assert type(result[0]).__name__ == "AIMessage"
|
||||
assert result[0].id == "c1"
|
||||
assert result[1].id == "h1"
|
||||
|
||||
|
||||
def test_left_as_dicts_is_normalized():
|
||||
"""Opt-1 guard: dicts at left[0] must trigger full conversion."""
|
||||
left = [{"role": "user", "content": "hi", "id": "d1"}]
|
||||
right = [AIMessage(content="reply", id="a1")]
|
||||
result = add_messages(left, right)
|
||||
assert len(result) == 2
|
||||
assert isinstance(result[0], HumanMessage)
|
||||
assert result[0].id == "d1"
|
||||
assert result[0].content == "hi"
|
||||
|
||||
|
||||
def test_left_as_tuples_is_normalized():
|
||||
"""Opt-1 guard: tuple-form messages must trigger full conversion."""
|
||||
left = [("user", "hi")]
|
||||
right = [AIMessage(content="reply", id="a1")]
|
||||
result = add_messages(left, right)
|
||||
assert len(result) == 2
|
||||
assert isinstance(result[0], HumanMessage)
|
||||
# id is auto-assigned
|
||||
assert isinstance(result[0].id, str) and UUID(result[0].id, version=4)
|
||||
|
||||
|
||||
def test_left_first_msg_missing_id_is_normalized():
|
||||
"""Opt-1 guard: a BaseMessage without an id at left[0] falls to the else branch."""
|
||||
left = [HumanMessage(content="hi")] # no id
|
||||
right = [AIMessage(content="reply", id="a1")]
|
||||
result = add_messages(left, right)
|
||||
assert len(result) == 2
|
||||
# left's id must have been auto-assigned
|
||||
assert isinstance(result[0].id, str) and UUID(result[0].id, version=4)
|
||||
|
||||
|
||||
def test_duplicate_ids_in_right_with_nonempty_left():
|
||||
"""Opt-2 guard: intra-right duplicate ids must take slow path (dedup kept)."""
|
||||
left = [HumanMessage(content="prior", id="1")]
|
||||
right = [
|
||||
AIMessage(content="first", id="2"),
|
||||
AIMessage(content="second", id="2"),
|
||||
]
|
||||
result = add_messages(left, right)
|
||||
assert len(result) == 2
|
||||
assert result[0].id == "1"
|
||||
assert result[1].id == "2"
|
||||
assert result[1].content == "second"
|
||||
|
||||
|
||||
def test_right_with_none_ids_pure_append():
|
||||
"""Fast path still correct when right entries start with id=None (fresh uuids assigned)."""
|
||||
left = [HumanMessage(content="prior", id="1")]
|
||||
right = [AIMessage(content="a"), AIMessage(content="b")]
|
||||
result = add_messages(left, right)
|
||||
assert len(result) == 3
|
||||
assert result[0].id == "1"
|
||||
for m in result[1:]:
|
||||
assert isinstance(m.id, str) and UUID(m.id, version=4)
|
||||
# fresh uuids must be distinct
|
||||
assert result[1].id != result[2].id
|
||||
|
||||
|
||||
def test_fast_path_returns_fresh_list():
|
||||
"""Fast path must return a new list object (not mutate or alias left)."""
|
||||
left = [HumanMessage(content="prior", id="1")]
|
||||
right = [AIMessage(content="new", id="2")]
|
||||
result = add_messages(left, right)
|
||||
assert result is not left
|
||||
# left must be untouched
|
||||
assert len(left) == 1
|
||||
assert left[0].id == "1"
|
||||
|
||||
|
||||
def test_push_messages_in_graph():
|
||||
class MessagesState(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
|
||||
@@ -9400,190 +9400,3 @@ def test_fork_does_not_apply_pending_writes(
|
||||
|
||||
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
|
||||
assert result == {"value": 121}
|
||||
|
||||
|
||||
async def test_delta_channel_end_to_end_inmemory() -> None:
|
||||
"""Full graph run: DeltaChannel accumulates correctly across multiple turns."""
|
||||
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
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-test-1"}}
|
||||
|
||||
# Turn 1
|
||||
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
||||
# Turn 2
|
||||
graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config)
|
||||
# Turn 3
|
||||
graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# 3 human + 3 AI = 6 total
|
||||
assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].content == "hello"
|
||||
assert msgs[2].content == "world"
|
||||
assert msgs[4].content == "bye"
|
||||
assert msgs[1].content == "reply-1"
|
||||
assert msgs[3].content == "reply-3"
|
||||
assert msgs[5].content == "reply-5"
|
||||
|
||||
|
||||
async def test_delta_channel_time_travel() -> None:
|
||||
"""Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas."""
|
||||
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
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
counter["n"] += 1
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
|
||||
]
|
||||
}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-time-travel"}}
|
||||
|
||||
# Run 2 turns: h1→ai-1, h2→ai-2
|
||||
graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config)
|
||||
graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config)
|
||||
|
||||
# Find the checkpoint after turn 1 (2 messages: h1 + ai-1)
|
||||
history = list(graph.get_state_history(config))
|
||||
after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2)
|
||||
|
||||
assert len(after_turn1.values["messages"]) == 2
|
||||
assert after_turn1.values["messages"][0].content == "h1"
|
||||
assert after_turn1.values["messages"][1].content == "ai-1"
|
||||
|
||||
# Resume from turn-1 checkpoint: inject h3, expect 3 messages total (h1, ai-1, ai-N)
|
||||
# NOT 5 messages (turn-2 deltas must not bleed into the resumed run)
|
||||
result = graph.invoke(
|
||||
{"messages": [HumanMessage(content="h3", id="h3")]},
|
||||
after_turn1.config,
|
||||
)
|
||||
msgs = result["messages"]
|
||||
# Should be: h1, ai-1, h3, ai-N — 4 messages total
|
||||
assert len(msgs) == 4, (
|
||||
f"expected 4 messages after time-travel resume, got {len(msgs)}: {msgs}"
|
||||
)
|
||||
assert msgs[0].content == "h1"
|
||||
assert msgs[1].content == "ai-1"
|
||||
assert msgs[2].content == "h3"
|
||||
|
||||
|
||||
async def test_delta_channel_remove_message_end_to_end() -> None:
|
||||
"""RemoveMessage inside a DeltaChannel graph must persist and reload correctly."""
|
||||
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
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
return {"messages": [AIMessage(content="reply", id="ai-1")]}
|
||||
|
||||
def delete_first(state: State) -> dict:
|
||||
# removes the first message
|
||||
return {"messages": [RemoveMessage(id=state["messages"][0].id)]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_node("delete_first", delete_first)
|
||||
builder.add_edge(START, "respond")
|
||||
builder.add_edge("respond", "delete_first")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-remove-test"}}
|
||||
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# h1 was removed, only ai-1 should remain
|
||||
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].id == "ai-1"
|
||||
|
||||
# A subsequent turn must reconstruct from the checkpoint correctly
|
||||
graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config)
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed
|
||||
# more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2
|
||||
# just assert h1 is still gone
|
||||
assert all(m.id != "h1" for m in msgs), (
|
||||
"h1 should still be absent after second turn"
|
||||
)
|
||||
|
||||
|
||||
async def test_delta_channel_update_by_id_end_to_end() -> None:
|
||||
"""Updating a message by ID via DeltaChannel must persist and reload correctly."""
|
||||
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
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def update_msg(state: State) -> dict:
|
||||
# re-send h1 with updated content
|
||||
return {"messages": [HumanMessage(content="updated", id="h1")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("update_msg", update_msg)
|
||||
builder.add_edge(START, "update_msg")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-update-id-test"}}
|
||||
graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].content == "updated"
|
||||
assert msgs[0].id == "h1"
|
||||
|
||||
# Second turn: verify the updated state is the base for further accumulation
|
||||
graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config)
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
ids = [m.id for m in msgs]
|
||||
assert "h1" in ids # h1 persists (updated, not duplicated)
|
||||
assert "h2" in ids
|
||||
assert ids.count("h1") == 1, "h1 must not be duplicated"
|
||||
|
||||
@@ -1,185 +0,0 @@
|
||||
"""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,6 +1113,70 @@ 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.0a2"
|
||||
version = "1.3.0"
|
||||
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/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
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" }
|
||||
wheels = [
|
||||
{ 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" },
|
||||
{ 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" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.7a2"
|
||||
version = "1.1.9"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1439,7 +1439,7 @@ test = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ 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.8.0" },
|
||||
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.9.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.9"
|
||||
version = "1.0.10"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -614,6 +614,7 @@ class _InjectedArgs:
|
||||
store: str | None
|
||||
runtime: str | None
|
||||
all_injected_keys: set[str]
|
||||
_optional_state_args: set[str]
|
||||
|
||||
|
||||
class ToolNode(RunnableCallable):
|
||||
@@ -807,6 +808,7 @@ 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,
|
||||
)
|
||||
@@ -841,6 +843,7 @@ 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,
|
||||
)
|
||||
@@ -1333,7 +1336,7 @@ class ToolNode(RunnableCallable):
|
||||
return tool_call
|
||||
|
||||
tool_call_copy: ToolCall = copy(tool_call)
|
||||
injected_args = {}
|
||||
injected_args: dict[str, Any] = {}
|
||||
|
||||
# Inject state
|
||||
if injected.state:
|
||||
@@ -1361,14 +1364,20 @@ class ToolNode(RunnableCallable):
|
||||
# Extract state values
|
||||
if isinstance(state, dict):
|
||||
for tool_arg, state_field in injected.state.items():
|
||||
injected_args[tool_arg] = (
|
||||
state[state_field] if state_field else state
|
||||
)
|
||||
if not state_field:
|
||||
injected_args[tool_arg] = state
|
||||
elif state_field in state:
|
||||
injected_args[tool_arg] = state[state_field]
|
||||
elif tool_arg not in injected._optional_state_args:
|
||||
raise KeyError(state_field)
|
||||
else:
|
||||
for tool_arg, state_field in injected.state.items():
|
||||
injected_args[tool_arg] = (
|
||||
getattr(state, state_field) if state_field else state
|
||||
)
|
||||
if not state_field:
|
||||
injected_args[tool_arg] = state
|
||||
elif hasattr(state, state_field):
|
||||
injected_args[tool_arg] = getattr(state, state_field)
|
||||
elif tool_arg not in injected._optional_state_args:
|
||||
raise AttributeError(state_field)
|
||||
|
||||
# Inject store
|
||||
if injected.store:
|
||||
@@ -1569,6 +1578,7 @@ 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.
|
||||
@@ -1611,6 +1621,7 @@ 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
|
||||
@@ -1859,6 +1870,7 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
store_arg: str | None = None
|
||||
runtime_arg: str | None = None
|
||||
all_injected_keys: set[str] = set()
|
||||
_optional_state_args: set[str] = set()
|
||||
|
||||
for name, type_ in all_annotations.items():
|
||||
# Track all InjectedToolArg-annotated params (including custom subclasses)
|
||||
@@ -1873,6 +1885,9 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
if state_inj := _get_injection_from_type(type_, InjectedState):
|
||||
if isinstance(state_inj, InjectedState) and state_inj.field:
|
||||
state_args[name] = state_inj.field
|
||||
field_info = full_schema.model_fields.get(name)
|
||||
if field_info and not field_info.is_required():
|
||||
_optional_state_args.add(name)
|
||||
else:
|
||||
state_args[name] = None
|
||||
|
||||
@@ -1889,4 +1904,5 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
store=store_arg,
|
||||
runtime=runtime_arg,
|
||||
all_injected_keys=all_injected_keys,
|
||||
_optional_state_args=_optional_state_args,
|
||||
)
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.9"
|
||||
version = "1.0.10"
|
||||
description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -69,6 +69,7 @@ def _create_config_with_runtime(store=None, state=None):
|
||||
context={},
|
||||
store=store,
|
||||
stream_writer=None,
|
||||
tools=[],
|
||||
tool_call_id="test_id",
|
||||
)
|
||||
return {
|
||||
|
||||
@@ -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_and_server_info() -> None:
|
||||
"""Test that execution_info and server_info are forwarded from Runtime to ToolRuntime."""
|
||||
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."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
|
||||
exec_info = ExecutionInfo(
|
||||
@@ -2043,9 +2043,15 @@ def test_tool_runtime_forwards_execution_info_and_server_info() -> None:
|
||||
"""Tool that captures runtime info."""
|
||||
captured["execution_info"] = runtime.execution_info
|
||||
captured["server_info"] = runtime.server_info
|
||||
captured["tools"] = runtime.tools
|
||||
return "ok"
|
||||
|
||||
node = ToolNode([info_tool])
|
||||
@dec_tool
|
||||
def other_tool(y: int) -> str:
|
||||
"""Another tool available to the runtime."""
|
||||
return str(y)
|
||||
|
||||
node = ToolNode([info_tool, other_tool])
|
||||
tool_call = {
|
||||
"name": "info_tool",
|
||||
"args": {"x": 1},
|
||||
@@ -2054,17 +2060,21 @@ def test_tool_runtime_forwards_execution_info_and_server_info() -> None:
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
||||
node.invoke({"messages": [msg]}, config=config)
|
||||
result = 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_and_server_info_async() -> None:
|
||||
"""Test that execution_info and server_info are forwarded in async path."""
|
||||
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."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
|
||||
exec_info = ExecutionInfo(
|
||||
@@ -2090,9 +2100,15 @@ async def test_tool_runtime_forwards_execution_info_and_server_info_async() -> N
|
||||
"""Async tool that captures runtime info."""
|
||||
captured["execution_info"] = runtime.execution_info
|
||||
captured["server_info"] = runtime.server_info
|
||||
captured["tools"] = runtime.tools
|
||||
return "ok"
|
||||
|
||||
node = ToolNode([info_tool_async])
|
||||
@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])
|
||||
tool_call = {
|
||||
"name": "info_tool_async",
|
||||
"args": {"x": 1},
|
||||
@@ -2101,12 +2117,17 @@ async def test_tool_runtime_forwards_execution_info_and_server_info_async() -> N
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
||||
await node.ainvoke({"messages": [msg]}, config=config)
|
||||
result = 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.0a2"
|
||||
version = "1.3.0"
|
||||
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/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
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" }
|
||||
wheels = [
|
||||
{ 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" },
|
||||
{ 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" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.7a2"
|
||||
version = "1.1.9"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -281,7 +281,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ 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.9"
|
||||
version = "1.0.10"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -1,30 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph_sdk.auth import Auth
|
||||
from langgraph_sdk.client import get_client, get_sync_client
|
||||
from langgraph_sdk.encryption import Encryption
|
||||
from langgraph_sdk.encryption.types import EncryptionContext
|
||||
from langgraph_sdk.auth import Auth
|
||||
from langgraph_sdk.client import get_client, get_sync_client
|
||||
from langgraph_sdk.encryption import Encryption
|
||||
from langgraph_sdk.encryption.types import EncryptionContext
|
||||
|
||||
__version__ = "0.3.13"
|
||||
|
||||
__all__ = ["Auth", "Encryption", "EncryptionContext", "get_client", "get_sync_client"]
|
||||
|
||||
_LAZY: dict[str, str] = {
|
||||
"Auth": "langgraph_sdk.auth",
|
||||
"get_client": "langgraph_sdk.client",
|
||||
"get_sync_client": "langgraph_sdk.client",
|
||||
"Encryption": "langgraph_sdk.encryption",
|
||||
"EncryptionContext": "langgraph_sdk.encryption.types",
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> object:
|
||||
if name in _LAZY:
|
||||
mod = importlib.import_module(_LAZY[name])
|
||||
return getattr(mod, name)
|
||||
msg = f"module {__name__!r} has no attribute {name!r}"
|
||||
raise AttributeError(msg)
|
||||
|
||||
Generated
+6
-6
@@ -262,7 +262,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.0a2"
|
||||
version = "1.3.0"
|
||||
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/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
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" }
|
||||
wheels = [
|
||||
{ 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" },
|
||||
{ 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" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.7a2"
|
||||
version = "1.1.9"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -294,7 +294,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ 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.9"
|
||||
version = "1.0.10"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Reference in New Issue
Block a user