mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-27 01:52:25 +02:00
Compare commits
40
Commits
@@ -121,8 +121,8 @@ jobs:
|
||||
exit 1
|
||||
fi
|
||||
LANGCHAIN_OPENAI_VERSION=$(docker run --rm --entrypoint "" langgraph-test-h python -c "import sys; from importlib.metadata import version; v = version('langchain-openai'); print(v);")
|
||||
if [ "$LANGCHAIN_OPENAI_VERSION" != "1.1.14" ]; then
|
||||
echo "LANGCHAIN_OPENAI_VERSION != 1.1.14; $LANGCHAIN_OPENAI_VERSION"
|
||||
if [ "$LANGCHAIN_OPENAI_VERSION" != "1.0.1" ]; then
|
||||
echo "LANGCHAIN_OPENAI_VERSION != 1.0.1; $LANGCHAIN_OPENAI_VERSION"
|
||||
exit 1
|
||||
fi
|
||||
LANGCHAIN_ANTHROPIC_VERSION=$(docker run --rm --entrypoint "" langgraph-test-h python -c "import sys; from importlib.metadata import version; v = version('langchain-anthropic'); print(v);")
|
||||
|
||||
@@ -100,3 +100,4 @@ dmypy.json
|
||||
.turbo
|
||||
.editorconfig
|
||||
.scratch
|
||||
.worktrees/
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,405 @@
|
||||
# DiffChannel: Incremental Checkpoint Storage for Append-Style Reducers
|
||||
|
||||
**Date:** 2026-04-17
|
||||
**Status:** Approved for implementation
|
||||
**Scope:** `libs/checkpoint`, `libs/langgraph`, `libs/checkpoint-postgres`
|
||||
|
||||
---
|
||||
|
||||
## Motivation
|
||||
|
||||
LangGraph checkpoints today store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, this means each checkpoint blob contains the entire conversation history. Storage cost grows O(N²) in the number of turns: step 1 stores 1 message, step 100 stores 100 messages, step 1000 stores 1000 messages. For long-running agentic conversations with high-token messages this is untenable.
|
||||
|
||||
The fix is to store only the **delta** (new writes) per step, reconstructing the full accumulated value at load time by replaying the chain. This is an opt-in mechanism — existing graphs are unaffected.
|
||||
|
||||
---
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- **Compaction / materialized snapshots**: deferred. Load cost stays O(N) blob fetches but those fetches are batched into a single query — acceptable for now.
|
||||
- **SQLite saver support**: SQLite stores all channel values inline in one row (no per-channel blob table). Deferred to a follow-up.
|
||||
- **Automatic migration** of existing `BinaryOperatorAggregate` channels: users opt in explicitly. Old checkpoints load correctly via the backwards-compatibility path in `from_checkpoint`.
|
||||
|
||||
---
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
```
|
||||
User state definition
|
||||
└── Annotated[list[AnyMessage], DiffChannel(add_messages)]
|
||||
|
||||
Write path (per superstep)
|
||||
DiffChannel.update() — apply operator, accumulate writes in _pending
|
||||
DiffChannel.checkpoint() — return DiffDelta(delta=_pending, prev_version=_base_version)
|
||||
serde.dumps_typed() — serialize DiffDelta as ("diff", msgpack_bytes)
|
||||
saver.put() — store blob at (thread_id, ns, "messages", version_N)
|
||||
DiffChannel.after_checkpoint(version_N) — advance _base_version, clear _pending
|
||||
|
||||
Read path (on graph load or time-travel)
|
||||
saver.get_tuple() — fetch current-version blob per channel
|
||||
saver._load_blobs() — detect "diff" type → follow chain to reconstruct DiffChainValue
|
||||
DiffChannel.from_checkpoint(DiffChainValue) — replay deltas with operator → full list
|
||||
DiffChannel.after_checkpoint(version_N) — set _base_version for next write
|
||||
```
|
||||
|
||||
The pregel layer (`_checkpoint.py`, `_loop.py`) is unchanged except for two small additions to call the new `after_checkpoint` hook. The saver public interface (`BaseCheckpointSaver`) gains no new methods. All chain-following logic lives inside each saver's private `_load_blobs`.
|
||||
|
||||
---
|
||||
|
||||
## New Protocol Types
|
||||
|
||||
**Location:** `libs/checkpoint/langgraph/checkpoint/base/__init__.py`
|
||||
|
||||
Two dataclasses form the contract between `DiffChannel` and savers:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class DiffDelta:
|
||||
"""Returned by DiffChannel.checkpoint(). Written to the blob store."""
|
||||
delta: list[Any] # raw writes passed to update() this step
|
||||
prev_version: str | None # version of the previous diff blob; None = chain root
|
||||
```
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class DiffChainValue:
|
||||
"""Passed to DiffChannel.from_checkpoint(). Assembled by _load_blobs()."""
|
||||
base: list[Any] | None # starting accumulated value (None = empty start)
|
||||
deltas: list[list[Any]] # write-sets ordered oldest → newest
|
||||
```
|
||||
|
||||
`DiffDelta` lives in the checkpoint base package (not the channel module) so savers can import it without creating a circular dependency. `DiffChainValue` is there for the same reason.
|
||||
|
||||
---
|
||||
|
||||
## `BaseChannel.after_checkpoint()` Hook
|
||||
|
||||
**Location:** `libs/langgraph/langgraph/channels/base.py`
|
||||
|
||||
```python
|
||||
def after_checkpoint(self, version: Any) -> None:
|
||||
"""Called after checkpoint() (with the new version) and after from_checkpoint()
|
||||
(with the current version). No-op by default; DiffChannel overrides."""
|
||||
pass
|
||||
```
|
||||
|
||||
This is a **non-abstract, no-op default** — fully backwards compatible. All existing channels inherit it silently. It is NOT in the abstract interface.
|
||||
|
||||
---
|
||||
|
||||
## `DiffChannel[V]`
|
||||
|
||||
**Location:** `libs/langgraph/langgraph/channels/diff.py` (new file)
|
||||
|
||||
### Internal state
|
||||
|
||||
| Attribute | Type | Description |
|
||||
|---|---|---|
|
||||
| `value` | `list[V]` | Full accumulated value (the reconstructed list) |
|
||||
| `operator` | `Callable` | The binary reducer (e.g. `add_messages`) |
|
||||
| `_pending` | `list[Any]` | Raw writes accumulated since last `after_checkpoint` call |
|
||||
| `_base_version` | `str \| None` | Version this channel was last checkpointed at (= `prev_version` for next delta) |
|
||||
| `_overwritten` | `bool` | True if an `Overwrite` was applied since last `after_checkpoint`; makes next blob a chain root |
|
||||
|
||||
### `update(values)`
|
||||
|
||||
Mirrors `BinaryOperatorAggregate.update()` with two additions:
|
||||
|
||||
1. For each non-Overwrite value: apply `self.operator(self.value, value)` as before; **also append the raw incoming value to `self._pending`**.
|
||||
2. For an `Overwrite(v)` value: set `self.value = v`; set `self._pending = list(v)` (full value becomes the new delta); set `self._overwritten = True`.
|
||||
|
||||
The key: `_pending` stores the **incoming writes** (what was passed to `update()`), not the diff of `self.value`. This is important because `add_messages` handles removal and update-by-ID — replaying the writes with `operator` during reconstruction applies that logic correctly.
|
||||
|
||||
### `checkpoint()`
|
||||
|
||||
```python
|
||||
def checkpoint(self) -> DiffDelta:
|
||||
return DiffDelta(
|
||||
delta=self._pending[:],
|
||||
prev_version=None if self._overwritten else self._base_version,
|
||||
)
|
||||
```
|
||||
|
||||
- Normal step: `prev_version = self._base_version` → chain link
|
||||
- After Overwrite: `prev_version = None` → chain root (reconstruction stops here and uses `delta` as the full base value)
|
||||
|
||||
Returns `DiffDelta`, never the raw accumulated list. The serde handles serialization.
|
||||
|
||||
### `from_checkpoint(checkpoint)`
|
||||
|
||||
```python
|
||||
def from_checkpoint(self, checkpoint) -> Self:
|
||||
new = DiffChannel(self.typ, self.operator)
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
new.value = []
|
||||
elif isinstance(checkpoint, DiffChainValue):
|
||||
accumulated = checkpoint.base or []
|
||||
for step_writes in checkpoint.deltas:
|
||||
# Mirror update() exactly: apply each write individually so operator
|
||||
# semantics (e.g. add_messages ID-based removal) are respected.
|
||||
for write in step_writes:
|
||||
accumulated = new.operator(accumulated, write)
|
||||
new.value = accumulated
|
||||
elif isinstance(checkpoint, DiffDelta):
|
||||
# Unsupported saver: _load_blobs returned a raw DiffDelta instead of
|
||||
# assembling a DiffChainValue. Raise rather than silently losing history.
|
||||
raise ValueError(
|
||||
"DiffChannel received a raw DiffDelta from the checkpoint saver. "
|
||||
"Your saver does not support incremental channel storage. "
|
||||
"Use InMemorySaver or PostgresSaver."
|
||||
)
|
||||
else:
|
||||
# Backwards compat: plain list from old BinaryOperatorAggregate checkpoint.
|
||||
new.value = checkpoint
|
||||
new._pending = []
|
||||
new._base_version = None # set by the subsequent after_checkpoint() call
|
||||
return new
|
||||
```
|
||||
|
||||
The operator is available on `self` (the channel spec) so reconstruction is correct for any reducer — the saver never needs to know about `add_messages`.
|
||||
|
||||
`_pending` stores **individual writes** (each `value` from `update()`'s `values` sequence), so each `step_writes` list in `DiffChainValue.deltas` is replayed write-by-write — identical to the `update()` loop.
|
||||
|
||||
### `after_checkpoint(version)`
|
||||
|
||||
```python
|
||||
def after_checkpoint(self, version: Any) -> None:
|
||||
if version != self._base_version:
|
||||
self._base_version = version
|
||||
self._pending = []
|
||||
self._overwritten = False
|
||||
```
|
||||
|
||||
No-op when `version == self._base_version` (channel wasn't updated this step — blob was not written). Clears `_pending` and advances `_base_version` when the channel was actually checkpointed.
|
||||
|
||||
### Opt-in API
|
||||
|
||||
```python
|
||||
from langgraph.channels.diff import DiffChannel
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DiffChannel(add_messages)]
|
||||
```
|
||||
|
||||
`StateGraph` already handles `BaseChannel` instances as annotation metadata — `DiffChannel` inherits this without any changes to `StateGraph`.
|
||||
|
||||
---
|
||||
|
||||
## Serde Extension
|
||||
|
||||
**Location:** `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py`
|
||||
|
||||
Add one branch to `dumps_typed` (before the `else` msgpack fallback), using the existing module-level `_msgpack_enc` so message ext-types (Pydantic v2, etc.) are handled correctly:
|
||||
|
||||
```python
|
||||
elif isinstance(obj, DiffDelta):
|
||||
return "diff", _msgpack_enc({"d": obj.delta, "p": obj.prev_version})
|
||||
```
|
||||
|
||||
Add one branch to `loads_typed` so savers can decode diff blobs without importing `ormsgpack` directly:
|
||||
|
||||
```python
|
||||
elif type_ == "diff":
|
||||
return ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
# returns {"d": [writes...], "p": prev_version_str_or_none}
|
||||
```
|
||||
|
||||
Savers call `serde.loads_typed(("diff", raw_bytes))` to decode a diff blob into `{"d": ..., "p": ...}`, then check `type_tag == "diff"` to trigger chain traversal. The serde layer is the only place that knows about `ormsgpack`.
|
||||
|
||||
---
|
||||
|
||||
## Saver Changes
|
||||
|
||||
### InMemorySaver
|
||||
|
||||
**`put()` — `libs/checkpoint/langgraph/checkpoint/memory/__init__.py`**
|
||||
|
||||
No change needed. The existing `self.serde.dumps_typed(values[k])` call already handles `DiffDelta` via the new serde branch above, storing it as `("diff", bytes)`.
|
||||
|
||||
**`_load_blobs()` — same file**
|
||||
|
||||
After checking `vv[0] != "empty"`, add a branch for `"diff"` before calling `serde.loads_typed`:
|
||||
|
||||
```python
|
||||
def _load_blobs(self, thread_id, checkpoint_ns, versions):
|
||||
channel_values = {}
|
||||
diff_channels = {} # channel_name -> current_version for diff channels
|
||||
|
||||
for k, v in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, v)
|
||||
if kk not in self.blobs:
|
||||
continue
|
||||
type_tag, blob_bytes = self.blobs[kk]
|
||||
if type_tag == "diff":
|
||||
diff_channels[k] = v # handle below
|
||||
elif type_tag != "empty":
|
||||
channel_values[k] = self.serde.loads_typed((type_tag, blob_bytes))
|
||||
|
||||
for k, current_version in diff_channels.items():
|
||||
# Follow chain: newest → oldest, then reverse
|
||||
chain_deltas = []
|
||||
base = None
|
||||
version = current_version
|
||||
while version is not None:
|
||||
kk = (thread_id, checkpoint_ns, k, version)
|
||||
if kk not in self.blobs:
|
||||
break
|
||||
type_tag, blob_bytes = self.blobs[kk]
|
||||
if type_tag == "diff":
|
||||
# Use serde so we don't need to import ormsgpack directly
|
||||
payload = self.serde.loads_typed((type_tag, blob_bytes))
|
||||
chain_deltas.append(payload["d"])
|
||||
version = payload["p"] # prev_version; None = root
|
||||
else:
|
||||
# Old non-diff blob encountered: treat as base accumulated value
|
||||
base = self.serde.loads_typed((type_tag, blob_bytes))
|
||||
break
|
||||
chain_deltas.reverse()
|
||||
channel_values[k] = DiffChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
return channel_values
|
||||
```
|
||||
|
||||
Each blob lookup is O(1) on the dict. Total: N dict lookups for a chain of depth N. Memory usage is identical to loading a single full-list blob (same total bytes, split across N entries).
|
||||
|
||||
### PostgresSaver
|
||||
|
||||
**`_load_blobs()` — `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py`**
|
||||
|
||||
The existing `SELECT_SQL` fetches one blob per channel via a JOIN. After running that query, detect any `"diff"` channels in the result and issue one additional range query:
|
||||
|
||||
```python
|
||||
def _load_blobs(self, blob_values):
|
||||
if not blob_values:
|
||||
return {}
|
||||
|
||||
result = {}
|
||||
diff_channels = {} # channel_name -> current_version (as str)
|
||||
|
||||
for k, t, v in blob_values:
|
||||
channel = k.decode()
|
||||
type_tag = t.decode()
|
||||
if type_tag == "diff":
|
||||
# Decode via serde — no direct ormsgpack import needed
|
||||
payload = self.serde.loads_typed((type_tag, v))
|
||||
diff_channels[channel] = payload # store for chain fetch
|
||||
elif type_tag != "empty":
|
||||
result[channel] = self.serde.loads_typed((type_tag, v))
|
||||
|
||||
if diff_channels:
|
||||
result.update(self._load_diff_chains(diff_channels))
|
||||
|
||||
return result
|
||||
```
|
||||
|
||||
`_load_diff_chains` issues one SQL query per diff channel (typically just `messages`):
|
||||
|
||||
```sql
|
||||
SELECT version, type, blob
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = %s
|
||||
AND checkpoint_ns = %s
|
||||
AND channel = %s
|
||||
AND version <= %s
|
||||
ORDER BY version ASC
|
||||
```
|
||||
|
||||
In Python, iterate rows in ascending version order: if `type = "diff"`, accumulate the delta; if any other type is encountered, treat it as the base accumulated value and stop. Return `DiffChainValue(base=..., deltas=[...])`.
|
||||
|
||||
This results in **at most 2 queries total** for a graph with one `DiffChannel` — existing behaviour for all other channels is unchanged.
|
||||
|
||||
**`put()` / `_dump_blobs()`**
|
||||
|
||||
No change needed. `_dump_blobs` calls `self.serde.dumps_typed(v)` for each channel value in `new_versions`. When `v` is a `DiffDelta`, the serde produces `("diff", bytes)` which is stored as `type = "diff"` in `checkpoint_blobs`. The `ON CONFLICT DO NOTHING` semantics are preserved.
|
||||
|
||||
### SQLite
|
||||
|
||||
Deferred. `SqliteSaver` stores the entire checkpoint as a single serialized row — it has no per-channel blob table. Supporting `DiffChannel` on SQLite would require adding a new blobs table, which is a separate migration tracked separately.
|
||||
|
||||
---
|
||||
|
||||
## Pregel Layer Changes
|
||||
|
||||
### `channels_from_checkpoint` — `libs/langgraph/langgraph/pregel/_checkpoint.py`
|
||||
|
||||
After constructing each channel from its checkpoint value, call `after_checkpoint` so the channel records its current version:
|
||||
|
||||
```python
|
||||
channels = {}
|
||||
for k, v in channel_specs.items():
|
||||
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
ch.after_checkpoint(checkpoint["channel_versions"].get(k))
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
```
|
||||
|
||||
Existing channels get the no-op `after_checkpoint`. `DiffChannel` uses it to set `_base_version`.
|
||||
|
||||
### `PregelLoop._put_checkpoint` — `libs/langgraph/langgraph/pregel/_loop.py`
|
||||
|
||||
After `create_checkpoint(self.checkpoint, self.channels, self.step, ...)` returns and `do_checkpoint is True` and `self.channels is not None`, iterate channels and notify:
|
||||
|
||||
```python
|
||||
if do_checkpoint and self.channels:
|
||||
for k, ch in self.channels.items():
|
||||
ch.after_checkpoint(self.checkpoint["channel_versions"].get(k))
|
||||
```
|
||||
|
||||
This is called after `create_checkpoint` updates `self.checkpoint["channel_versions"]`, so `get(k)` returns the new version for updated channels and the old version for unchanged ones. `DiffChannel.after_checkpoint` only clears `_pending` when `version != _base_version`, so unchanged channels are no-ops.
|
||||
|
||||
---
|
||||
|
||||
## Backwards Compatibility
|
||||
|
||||
| Scenario | Behaviour |
|
||||
|---|---|
|
||||
| Existing graph using `add_messages` (BinaryOperatorAggregate) | Unaffected — no code changes, no data migration |
|
||||
| New graph with `DiffChannel`, loading old checkpoint blobs | `from_checkpoint` receives a plain `list` → used directly as accumulated value |
|
||||
| `DiffChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported |
|
||||
| `DiffChannel` with `SqliteSaver` | `from_checkpoint` receives a raw `DiffDelta` (SqliteSaver stores channel_values inline), raises `ValueError` with a clear message pointing to supported savers |
|
||||
| Time-travel / fork to past checkpoint | Chain traversal uses the version at that checkpoint → reconstruction is correct |
|
||||
| `update_state` | Treated as a normal step: writes are deltas chained to history |
|
||||
| `Overwrite` value | Resets chain: next blob has `prev_version=None`; reconstruction starts fresh |
|
||||
|
||||
---
|
||||
|
||||
## Testing Strategy
|
||||
|
||||
1. **Unit tests for `DiffChannel`** (`libs/langgraph/tests/`):
|
||||
- `update` → `checkpoint` → `after_checkpoint` → `checkpoint` lifecycle (2 steps, verify delta isolation)
|
||||
- `from_checkpoint(DiffChainValue)` correctly replays multi-step chains using the operator
|
||||
- `from_checkpoint(plain_list)` backwards-compat path
|
||||
- `Overwrite` creates a root blob (`prev_version=None`) and reconstruction ignores prior chain
|
||||
- `after_checkpoint` no-ops when version is unchanged
|
||||
|
||||
2. **Integration tests with `InMemorySaver`** (`libs/langgraph/tests/`):
|
||||
- 10-step conversation: verify final loaded state equals full accumulated messages
|
||||
- Time-travel: fork to step 5, verify only messages 1–5 are present
|
||||
- Mixed graph: some channels `BinaryOperatorAggregate`, one `DiffChannel` — both reconstruct correctly
|
||||
|
||||
3. **Serde tests** (`libs/checkpoint/tests/`):
|
||||
- `DiffDelta` round-trips through `dumps_typed` / saver storage
|
||||
- Old `"msgpack"` blob for a channel → `DiffChannel.from_checkpoint` handles it
|
||||
|
||||
4. **Postgres integration tests** (`libs/checkpoint-postgres/tests/`):
|
||||
- Range query reconstructs correct full list after N steps
|
||||
- Time-travel to checkpoint M reconstructs correct list of M messages
|
||||
|
||||
---
|
||||
|
||||
## Files Changed
|
||||
|
||||
| File | Change |
|
||||
|---|---|
|
||||
| `libs/checkpoint/langgraph/checkpoint/base/__init__.py` | Add `DiffDelta`, `DiffChainValue` dataclasses |
|
||||
| `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` | Add `"diff"` branch in `dumps_typed` |
|
||||
| `libs/checkpoint/langgraph/checkpoint/memory/__init__.py` | Chain traversal in `_load_blobs` |
|
||||
| `libs/langgraph/langgraph/channels/base.py` | Add no-op `after_checkpoint` method |
|
||||
| `libs/langgraph/langgraph/channels/diff.py` | **New file** — `DiffChannel` implementation |
|
||||
| `libs/langgraph/langgraph/channels/__init__.py` | Export `DiffChannel` |
|
||||
| `libs/langgraph/langgraph/pregel/_checkpoint.py` | Call `after_checkpoint` in `channels_from_checkpoint` |
|
||||
| `libs/langgraph/langgraph/pregel/_loop.py` | Call `after_checkpoint` after `create_checkpoint` |
|
||||
| `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` | Range-query chain reconstruction in `_load_blobs` |
|
||||
@@ -430,6 +430,43 @@ 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.
|
||||
@@ -442,6 +479,13 @@ 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": {
|
||||
@@ -454,7 +498,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
**channel_values,
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
|
||||
@@ -391,6 +391,43 @@ 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.
|
||||
@@ -403,11 +440,19 @@ 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": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
@@ -415,15 +460,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
**channel_values,
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -185,15 +185,22 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
)
|
||||
|
||||
def _load_blobs(
|
||||
self, blob_values: list[tuple[bytes, bytes, bytes]]
|
||||
self,
|
||||
blob_values: list[tuple[bytes, bytes, bytes]],
|
||||
*,
|
||||
thread_id: str = "",
|
||||
checkpoint_ns: str = "",
|
||||
cur: Any = None,
|
||||
) -> dict[str, Any]:
|
||||
if not blob_values:
|
||||
return {}
|
||||
return {
|
||||
k.decode(): self.serde.loads_typed((t.decode(), v))
|
||||
for k, t, v in blob_values
|
||||
if t.decode() != "empty"
|
||||
}
|
||||
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
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
|
||||
@@ -371,3 +371,47 @@ 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"
|
||||
|
||||
Generated
+1
-1
@@ -259,7 +259,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Generated
+1
-1
@@ -268,7 +268,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import dataclasses
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
|
||||
from typing import ( # noqa: UP035
|
||||
@@ -28,7 +30,46 @@ 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 during checkpoint hydration."""
|
||||
|
||||
base: list[Any] | None # starting accumulated value; None = start from empty
|
||||
deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest
|
||||
|
||||
|
||||
CheckpointHydrationKind = Literal["delta"]
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class IncrementalChannelSpec:
|
||||
"""Describes a checkpoint field that needs saver-side materialization."""
|
||||
|
||||
name: str
|
||||
kind: CheckpointHydrationKind
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class CheckpointHydrationPlan:
|
||||
"""Lists the checkpoint fields eligible for saver-side materialization."""
|
||||
|
||||
channels: tuple[IncrementalChannelSpec, ...]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_MISSING_SENTINEL = object()
|
||||
|
||||
|
||||
# Marked as total=False to allow for future expansion.
|
||||
@@ -457,6 +498,172 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def materialize_checkpoint(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Materialize any saver-managed incremental values in a checkpoint."""
|
||||
if plan is None or not plan.channels:
|
||||
return checkpoint
|
||||
|
||||
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 spec in plan.channels:
|
||||
value = checkpoint["channel_values"].get(spec.name)
|
||||
if spec.kind != "delta" or not isinstance(value, DeltaValue):
|
||||
continue
|
||||
|
||||
assembled_value = self._materialize_delta_channel(
|
||||
thread_id=thread_id,
|
||||
checkpoint_ns=checkpoint_ns,
|
||||
current_checkpoint_id=current_checkpoint_id,
|
||||
channel=spec.name,
|
||||
value=value,
|
||||
)
|
||||
if assembled_value is not None:
|
||||
assembled[spec.name] = assembled_value
|
||||
|
||||
if not assembled:
|
||||
return checkpoint
|
||||
|
||||
return {
|
||||
**checkpoint,
|
||||
"channel_values": {**checkpoint["channel_values"], **assembled},
|
||||
}
|
||||
|
||||
def materialize_checkpoint_tuple(
|
||||
self,
|
||||
value: CheckpointTuple,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> CheckpointTuple:
|
||||
"""Materialize incremental values for a single checkpoint tuple."""
|
||||
checkpoint = self.materialize_checkpoint(value.config, value.checkpoint, plan)
|
||||
if checkpoint is value.checkpoint:
|
||||
return value
|
||||
return value._replace(checkpoint=checkpoint)
|
||||
|
||||
def materialize_checkpoint_tuples(
|
||||
self,
|
||||
values: Sequence[CheckpointTuple],
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Sequence[CheckpointTuple]:
|
||||
"""Materialize incremental values for a batch of checkpoint tuples."""
|
||||
return [self.materialize_checkpoint_tuple(value, plan) for value in values]
|
||||
|
||||
async def amaterialize_checkpoint(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Async materialization hook for saver-managed incremental values."""
|
||||
if plan is None or not plan.channels:
|
||||
return checkpoint
|
||||
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
current_checkpoint_id = checkpoint.get("id")
|
||||
|
||||
targets = [
|
||||
(spec, checkpoint["channel_values"][spec.name])
|
||||
for spec in plan.channels
|
||||
if spec.kind == "delta"
|
||||
and isinstance(checkpoint["channel_values"].get(spec.name), DeltaValue)
|
||||
]
|
||||
if not targets:
|
||||
return checkpoint
|
||||
|
||||
# Walks for independent channels can run concurrently — each has its own chain.
|
||||
results = await asyncio.gather(
|
||||
*(
|
||||
self._amaterialize_delta_channel(
|
||||
thread_id=thread_id,
|
||||
checkpoint_ns=checkpoint_ns,
|
||||
current_checkpoint_id=current_checkpoint_id,
|
||||
channel=spec.name,
|
||||
value=value,
|
||||
)
|
||||
for spec, value in targets
|
||||
)
|
||||
)
|
||||
assembled = {
|
||||
spec.name: result
|
||||
for (spec, _), result in zip(targets, results, strict=True)
|
||||
if result is not None
|
||||
}
|
||||
|
||||
if not assembled:
|
||||
return checkpoint
|
||||
|
||||
return {
|
||||
**checkpoint,
|
||||
"channel_values": {**checkpoint["channel_values"], **assembled},
|
||||
}
|
||||
|
||||
async def amaterialize_checkpoint_tuple(
|
||||
self,
|
||||
value: CheckpointTuple,
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> CheckpointTuple:
|
||||
"""Async materialization hook for a single checkpoint tuple."""
|
||||
checkpoint = await self.amaterialize_checkpoint(
|
||||
value.config, value.checkpoint, plan
|
||||
)
|
||||
if checkpoint is value.checkpoint:
|
||||
return value
|
||||
return value._replace(checkpoint=checkpoint)
|
||||
|
||||
async def amaterialize_checkpoint_tuples(
|
||||
self,
|
||||
values: Sequence[CheckpointTuple],
|
||||
plan: CheckpointHydrationPlan | None = None,
|
||||
) -> Sequence[CheckpointTuple]:
|
||||
"""Async materialization hook for a batch of checkpoint tuples."""
|
||||
return [
|
||||
await self.amaterialize_checkpoint_tuple(value, plan) for value in values
|
||||
]
|
||||
|
||||
def get_channel_blob(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> Any:
|
||||
"""Look up a single channel blob by checkpoint ID + channel name.
|
||||
|
||||
Returns NotImplemented if this saver does not support efficient
|
||||
per-channel-version blob lookup. The pregel layer will fall back to
|
||||
get_tuple() traversal in that case.
|
||||
|
||||
Savers with a dedicated blob store (InMemorySaver, PostgresSaver)
|
||||
should override this for O(1) performance.
|
||||
"""
|
||||
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.
|
||||
|
||||
@@ -489,6 +696,138 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
clone.serde = maybe_add_typed_methods(serde)
|
||||
return clone
|
||||
|
||||
def _materialize_delta_channel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
current_checkpoint_id: str | None,
|
||||
channel: str,
|
||||
value: DeltaValue,
|
||||
) -> DeltaChainValue | None:
|
||||
chain_deltas: list[list[Any]] = []
|
||||
base: list[Any] | None = None
|
||||
cursor: DeltaValue = value
|
||||
visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set()
|
||||
|
||||
while True:
|
||||
chain_deltas.append(cursor.delta)
|
||||
prev_id = cursor.prev_checkpoint_id
|
||||
if prev_id is None:
|
||||
break
|
||||
if prev_id in visited:
|
||||
logger.warning(
|
||||
"DeltaChannel chain cycle at checkpoint %r for channel %r; breaking",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
visited.add(prev_id)
|
||||
|
||||
blob = self.get_channel_blob(thread_id, checkpoint_ns, prev_id, channel)
|
||||
if blob is not NotImplemented:
|
||||
if isinstance(blob, DeltaValue):
|
||||
cursor = blob
|
||||
continue
|
||||
base = blob
|
||||
break
|
||||
|
||||
parent_config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": prev_id,
|
||||
}
|
||||
}
|
||||
parent_tuple = self.get_tuple(parent_config)
|
||||
if parent_tuple is None:
|
||||
logger.warning(
|
||||
"DeltaChannel chain broken: checkpoint %r not found for channel %r",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
prev_val = parent_tuple.checkpoint["channel_values"].get(
|
||||
channel, _MISSING_SENTINEL
|
||||
)
|
||||
if prev_val is _MISSING_SENTINEL:
|
||||
break
|
||||
if isinstance(prev_val, DeltaValue):
|
||||
cursor = prev_val
|
||||
else:
|
||||
base = prev_val
|
||||
break
|
||||
|
||||
chain_deltas.reverse()
|
||||
return DeltaChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
async def _amaterialize_delta_channel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
current_checkpoint_id: str | None,
|
||||
channel: str,
|
||||
value: DeltaValue,
|
||||
) -> DeltaChainValue | None:
|
||||
chain_deltas: list[list[Any]] = []
|
||||
base: list[Any] | None = None
|
||||
cursor: DeltaValue = value
|
||||
visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set()
|
||||
|
||||
while True:
|
||||
chain_deltas.append(cursor.delta)
|
||||
prev_id = cursor.prev_checkpoint_id
|
||||
if prev_id is None:
|
||||
break
|
||||
if prev_id in visited:
|
||||
logger.warning(
|
||||
"DeltaChannel chain cycle at checkpoint %r for channel %r; breaking",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
visited.add(prev_id)
|
||||
|
||||
blob = await self.aget_channel_blob(
|
||||
thread_id, checkpoint_ns, prev_id, channel
|
||||
)
|
||||
if blob is not NotImplemented:
|
||||
if isinstance(blob, DeltaValue):
|
||||
cursor = blob
|
||||
continue
|
||||
base = blob
|
||||
break
|
||||
|
||||
parent_config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": prev_id,
|
||||
}
|
||||
}
|
||||
parent_tuple = await self.aget_tuple(parent_config)
|
||||
if parent_tuple is None:
|
||||
logger.warning(
|
||||
"DeltaChannel chain broken: checkpoint %r not found for channel %r",
|
||||
prev_id,
|
||||
channel,
|
||||
)
|
||||
break
|
||||
prev_val = parent_tuple.checkpoint["channel_values"].get(
|
||||
channel, _MISSING_SENTINEL
|
||||
)
|
||||
if prev_val is _MISSING_SENTINEL:
|
||||
break
|
||||
if isinstance(prev_val, DeltaValue):
|
||||
cursor = prev_val
|
||||
else:
|
||||
base = prev_val
|
||||
break
|
||||
|
||||
chain_deltas.reverse()
|
||||
return DeltaChainValue(base=base, deltas=chain_deltas)
|
||||
|
||||
|
||||
def _with_msgpack_allowlist(
|
||||
serde: SerializerProtocol, extra_allowlist: Collection[tuple[str, ...]]
|
||||
|
||||
@@ -126,12 +126,46 @@ class InMemorySaver(
|
||||
channel_values: dict[str, Any] = {}
|
||||
for k, v in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, v)
|
||||
if kk in self.blobs:
|
||||
vv = self.blobs[kk]
|
||||
if vv[0] != "empty":
|
||||
channel_values[k] = self.serde.loads_typed(vv)
|
||||
if kk not in self.blobs:
|
||||
continue
|
||||
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,6 +80,8 @@ SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
|
||||
("langgraph.types", "Overwrite"),
|
||||
("langgraph.store.base", "Item"),
|
||||
("langgraph.store.base", "GetOp"),
|
||||
# DeltaChannel checkpoint value type
|
||||
("langgraph.checkpoint.base", "DeltaValue"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -46,35 +46,11 @@ LC_REVIVER = Reviver()
|
||||
EMPTY_BYTES = b""
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Dedup log warnings across process lifetime; cap bounds state if types are
|
||||
# dynamically generated (also acts as a circuit breaker on warning volume).
|
||||
# Dedup is best-effort: racing threads may each emit once for the same key,
|
||||
# and warnings are silently dropped once _MAX_WARNED_TYPES is reached.
|
||||
_MAX_WARNED_TYPES = 1000
|
||||
_warned_unregistered_types: set[tuple[str, str]] = set()
|
||||
_warned_blocked_types: set[tuple[str, str]] = set()
|
||||
|
||||
def _is_delta_value(obj: Any) -> bool:
|
||||
from langgraph.checkpoint.base import DeltaValue # lazy import avoids circular dep
|
||||
|
||||
def _is_safe_json_type(id_list: list[str]) -> bool:
|
||||
"""Return True if an lc=2 id refers to a type in SAFE_MSGPACK_TYPES.
|
||||
|
||||
Safe types bypass the ``allowed_json_modules`` gate so that old "json" format
|
||||
checkpoints (written before the msgpack migration) can be resumed without
|
||||
requiring users to configure an explicit allowlist.
|
||||
"""
|
||||
if len(id_list) < 2:
|
||||
return False
|
||||
module_name = ".".join(id_list[:-1])
|
||||
return (module_name, id_list[-1]) in _lg_msgpack.SAFE_MSGPACK_TYPES
|
||||
|
||||
|
||||
def _warn_once(
|
||||
seen: set[tuple[str, str]], key: tuple[str, str], msg: str, *args: object
|
||||
) -> None:
|
||||
if key in seen or len(seen) >= _MAX_WARNED_TYPES:
|
||||
return
|
||||
seen.add(key)
|
||||
logger.warning(msg, *args)
|
||||
return isinstance(obj, DeltaValue)
|
||||
|
||||
|
||||
class JsonPlusSerializer(SerializerProtocol):
|
||||
@@ -177,23 +153,19 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return out
|
||||
|
||||
def _reviver(self, value: dict[str, Any]) -> Any:
|
||||
if (
|
||||
if self._allowed_json_modules and (
|
||||
value.get("lc", None) == 2
|
||||
and value.get("type", None) == "constructor"
|
||||
and value.get("id", None) is not None
|
||||
):
|
||||
id_list = value["id"]
|
||||
is_safe = _is_safe_json_type(id_list)
|
||||
if self._allowed_json_modules or is_safe:
|
||||
try:
|
||||
return self._revive_lc2(value)
|
||||
except InvalidModuleError as e:
|
||||
if not is_safe:
|
||||
logger.warning(
|
||||
"Object %s is not in the deserialization allowlist.\n%s",
|
||||
value["id"],
|
||||
e.message,
|
||||
)
|
||||
try:
|
||||
return self._revive_lc2(value)
|
||||
except InvalidModuleError as e:
|
||||
logger.warning(
|
||||
"Object %s is not in the deserialization allowlist.\n%s",
|
||||
value["id"],
|
||||
e.message,
|
||||
)
|
||||
|
||||
return LC_REVIVER(value)
|
||||
|
||||
@@ -241,13 +213,6 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
method_display = "<init>"
|
||||
|
||||
dotted = ".".join(needed)
|
||||
# Safe types (the same set already allowed for msgpack deserialization) are
|
||||
# permitted without an explicit allowlist — they are known-safe LangGraph and
|
||||
# LangChain types. This restores backwards-compat for old "json" checkpoints
|
||||
# that pre-date the msgpack migration without reopening the broader security gate.
|
||||
if _is_safe_json_type(list(needed)):
|
||||
return
|
||||
|
||||
if not self._allowed_json_modules:
|
||||
raise InvalidModuleError(
|
||||
f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
|
||||
@@ -280,6 +245,8 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return "bytes", obj
|
||||
elif isinstance(obj, bytearray):
|
||||
return "bytearray", obj
|
||||
elif _is_delta_value(obj):
|
||||
return "delta", _msgpack_enc({"d": obj.delta, "c": obj.prev_checkpoint_id})
|
||||
else:
|
||||
try:
|
||||
return "msgpack", _msgpack_enc(obj)
|
||||
@@ -302,6 +269,13 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
elif type_ == "delta":
|
||||
from langgraph.checkpoint.base import DeltaValue # lazy import
|
||||
|
||||
raw = ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
return DeltaValue(delta=raw["d"], prev_checkpoint_id=raw.get("c"))
|
||||
elif self.pickle_fallback and type_ == "pickle":
|
||||
return pickle.loads(data_)
|
||||
else:
|
||||
@@ -575,9 +549,7 @@ def _create_msgpack_ext_hook(
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
_warn_once(
|
||||
_warned_unregistered_types,
|
||||
key,
|
||||
logger.warning(
|
||||
"Deserializing unregistered type %s.%s from checkpoint. "
|
||||
"This will be blocked in a future version. "
|
||||
"Set LANGGRAPH_STRICT_MSGPACK=true to block now, or add "
|
||||
@@ -599,9 +571,7 @@ def _create_msgpack_ext_hook(
|
||||
"name": name,
|
||||
}
|
||||
)
|
||||
_warn_once(
|
||||
_warned_blocked_types,
|
||||
key,
|
||||
logger.warning(
|
||||
"Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
|
||||
"Add to allowed_msgpack_modules to allow: [(%r, %r)]",
|
||||
module,
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -29,8 +29,6 @@ from langgraph.checkpoint.serde.jsonplus import (
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
|
||||
|
||||
@@ -104,13 +102,6 @@ def test_msgpack_method_pathlib_blocked_encrypted_strict(
|
||||
class TestEncryptedSerializerMsgpackAllowlist:
|
||||
"""Test msgpack allowlist behavior through EncryptedSerializer."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_warned_types(self) -> None:
|
||||
# Warning dedup state is process-global; reset per-test so each case
|
||||
# sees a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
|
||||
def test_safe_types_no_warning(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Test safe types deserialize without warnings through encryption."""
|
||||
serde = _make_encrypted_serde()
|
||||
|
||||
@@ -35,8 +35,6 @@ from langgraph.checkpoint.serde.jsonplus import (
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_msgpack_ext_hook_to_json,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
from langgraph.store.base import Item
|
||||
|
||||
@@ -333,57 +331,6 @@ def test_serde_jsonplus_bytes() -> None:
|
||||
assert serde.loads_typed(dumped) == some_bytes
|
||||
|
||||
|
||||
def test_lc2_json_safe_type_revives_without_allowlist() -> None:
|
||||
"""Old 'json' blobs with lc=2 for safe types must revive without an explicit allowlist.
|
||||
|
||||
Regression test for: https://github.com/langchain-ai/langgraph/issues/7498
|
||||
Threads checkpointed before v1.0.1 (pre-msgpack) stored messages as lc=2 JSON
|
||||
constructor dicts. Resuming those threads must reconstruct proper BaseMessage objects
|
||||
rather than returning raw dicts that cause MESSAGE_COERCION_FAILURE in add_messages.
|
||||
"""
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
serde = JsonPlusSerializer() # default: _allowed_json_modules=None
|
||||
|
||||
human_blob = {
|
||||
"lc": 2,
|
||||
"type": "constructor",
|
||||
"id": ["langchain_core", "messages", "human", "HumanMessage"],
|
||||
"kwargs": {"content": "hello", "type": "human"},
|
||||
}
|
||||
ai_blob = {
|
||||
"lc": 2,
|
||||
"type": "constructor",
|
||||
"id": ["langchain_core", "messages", "ai", "AIMessage"],
|
||||
"kwargs": {"content": "hi there", "type": "ai"},
|
||||
}
|
||||
result = serde.loads_typed(("json", json.dumps([human_blob, ai_blob]).encode()))
|
||||
|
||||
assert len(result) == 2
|
||||
assert isinstance(result[0], HumanMessage), (
|
||||
f"Expected HumanMessage, got {type(result[0])}: {result[0]!r}\n"
|
||||
"lc=2 JSON blobs for safe types must deserialize without an explicit allowlist"
|
||||
)
|
||||
assert result[0].content == "hello"
|
||||
assert isinstance(result[1], AIMessage)
|
||||
assert result[1].content == "hi there"
|
||||
|
||||
|
||||
def test_lc2_json_unknown_type_stays_blocked_without_allowlist() -> None:
|
||||
"""lc=2 JSON blobs for types NOT in SAFE_MSGPACK_TYPES still require an allowlist."""
|
||||
serde = JsonPlusSerializer()
|
||||
load = {
|
||||
"lc": 2,
|
||||
"type": "constructor",
|
||||
"id": ["pprint", "pprint"],
|
||||
"kwargs": {"object": "HELLO"},
|
||||
}
|
||||
# No allowlist configured → raw dict returned (not raised, not reconstructed)
|
||||
result = serde.loads_typed(("json", json.dumps(load).encode()))
|
||||
assert isinstance(result, dict), "Unknown lc=2 type must stay as raw dict"
|
||||
assert result.get("lc") == 2
|
||||
|
||||
|
||||
def test_deserde_invalid_module() -> None:
|
||||
serde = JsonPlusSerializer()
|
||||
load = {
|
||||
@@ -633,14 +580,6 @@ def test_msgpack_safe_types_no_warning(caplog: pytest.LogCaptureFixture) -> None
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_warned_types() -> None:
|
||||
# Warning dedup state is process-global; reset per-test so each case sees
|
||||
# a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
|
||||
|
||||
def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Pydantic models not in allowlist should log warning but still deserialize."""
|
||||
current = _lg_msgpack.STRICT_MSGPACK_ENABLED
|
||||
@@ -656,12 +595,6 @@ def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) ->
|
||||
assert "unregistered type" in caplog.text.lower()
|
||||
assert "allowed_msgpack_modules" in caplog.text
|
||||
assert result == obj
|
||||
|
||||
# Second deserialization of the same type should NOT produce another warning
|
||||
caplog.clear()
|
||||
result2 = serde.loads_typed(dumped)
|
||||
assert "unregistered type" not in caplog.text.lower()
|
||||
assert result2 == obj
|
||||
_lg_msgpack.STRICT_MSGPACK_ENABLED = current
|
||||
|
||||
|
||||
@@ -706,6 +639,7 @@ def test_msgpack_allowlist_silences_warning(caplog: pytest.LogCaptureFixture) ->
|
||||
|
||||
def test_msgpack_none_blocks_unregistered(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""allowed_msgpack_modules=None should block unregistered types."""
|
||||
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
|
||||
|
||||
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
|
||||
@@ -723,6 +657,7 @@ def test_msgpack_allowlist_blocks_non_listed(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""Allowlists should block unregistered types even if msgpack is enabled."""
|
||||
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")]
|
||||
)
|
||||
@@ -1048,3 +983,31 @@ 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,25 +12,13 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.jsonplus import (
|
||||
JsonPlusSerializer,
|
||||
_warned_blocked_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
|
||||
class MemoryPydantic(BaseModel):
|
||||
foo: str
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_warned_types() -> None:
|
||||
# Warning dedup state is process-global; reset per-test so each case sees
|
||||
# a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
|
||||
|
||||
class TestMemorySaver:
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self) -> None:
|
||||
@@ -320,3 +308,36 @@ 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
|
||||
)
|
||||
|
||||
Generated
+1
-1
@@ -286,7 +286,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -5,5 +5,5 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.1.14"
|
||||
"langchain-openai==1.0.1"
|
||||
]
|
||||
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.1.14",
|
||||
"langchain-openai==1.0.0a2",
|
||||
"langchain-anthropic==1.0.0a5",
|
||||
"langgraph==1.1.5"
|
||||
]
|
||||
|
||||
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain-openai==1.1.14",
|
||||
"langchain-openai==1.0.0a2",
|
||||
"langgraph==1.1.2",
|
||||
"langchain_community>=0.3.0",
|
||||
]
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.4.24"
|
||||
__version__ = "0.4.22"
|
||||
|
||||
@@ -1,124 +0,0 @@
|
||||
"""Shared ignore-file handling for local source filtering."""
|
||||
|
||||
import pathlib
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pathspec
|
||||
|
||||
_ALWAYS_EXCLUDE = [
|
||||
"__pycache__/",
|
||||
".git/",
|
||||
".venv/",
|
||||
"venv/",
|
||||
"node_modules/",
|
||||
".tox/",
|
||||
".mypy_cache/",
|
||||
]
|
||||
_ALWAYS_EXCLUDE_NAMES = frozenset(
|
||||
pattern.rstrip("/").split("/")[-1] for pattern in _ALWAYS_EXCLUDE
|
||||
)
|
||||
_GLOB_CHARS = frozenset("*?[")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _NegatedDockerignoreHints:
|
||||
exact_dirs: frozenset[pathlib.PurePosixPath] = frozenset()
|
||||
wildcard_prefixes: frozenset[pathlib.PurePosixPath] = frozenset()
|
||||
recurse_all: bool = False
|
||||
|
||||
def requires_dir_walk(self, path: pathlib.PurePosixPath) -> bool:
|
||||
if self.recurse_all or path in self.exact_dirs:
|
||||
return True
|
||||
return any(
|
||||
path == prefix or path in prefix.parents or prefix in path.parents
|
||||
for prefix in self.wildcard_prefixes
|
||||
)
|
||||
|
||||
|
||||
def _build_ignore_spec(
|
||||
directory: pathlib.Path, *, include_gitignore: bool = True
|
||||
) -> pathspec.PathSpec:
|
||||
"""Build a PathSpec combining built-in exclusions with ignore files.
|
||||
|
||||
Always excludes common non-source directories (`_ALWAYS_EXCLUDE`). On top
|
||||
of that, patterns from `.dockerignore` are merged in. `.gitignore` patterns
|
||||
are optional because some callers need Docker build-context semantics,
|
||||
while archive creation wants both files.
|
||||
"""
|
||||
lines: list[str] = list(_ALWAYS_EXCLUDE)
|
||||
ignore_files = [".dockerignore"]
|
||||
if include_gitignore:
|
||||
ignore_files.append(".gitignore")
|
||||
for name in ignore_files:
|
||||
ignore_file = directory / name
|
||||
if ignore_file.is_file():
|
||||
lines.extend(ignore_file.read_text(encoding="utf-8").splitlines())
|
||||
return pathspec.PathSpec.from_lines("gitwildmatch", lines)
|
||||
|
||||
|
||||
def _is_always_excluded(path: pathlib.PurePosixPath, *, is_dir: bool) -> bool:
|
||||
"""Whether `path` lives inside a built-in excluded directory."""
|
||||
parent_parts = path.parts if is_dir else path.parts[:-1]
|
||||
return any(part in _ALWAYS_EXCLUDE_NAMES for part in parent_parts)
|
||||
|
||||
|
||||
def _build_dockerignore_negation_hints(
|
||||
directory: pathlib.Path,
|
||||
) -> _NegatedDockerignoreHints:
|
||||
"""Summarize which ignored directories must still be traversed.
|
||||
|
||||
Most negations only require walking a small, concrete chain of parent
|
||||
directories (for example `!assets/keep.txt` requires entering `assets/`).
|
||||
Broader glob negations may force a wider walk.
|
||||
"""
|
||||
ignore_file = directory / ".dockerignore"
|
||||
if not ignore_file.is_file():
|
||||
return _NegatedDockerignoreHints()
|
||||
|
||||
exact_dirs: set[pathlib.PurePosixPath] = set()
|
||||
wildcard_prefixes: set[pathlib.PurePosixPath] = set()
|
||||
recurse_all = False
|
||||
|
||||
for raw_line in ignore_file.read_text(encoding="utf-8").splitlines():
|
||||
line = raw_line.strip()
|
||||
if not line or line.startswith("#") or line.startswith("\\!"):
|
||||
continue
|
||||
if line.startswith("\\#"):
|
||||
line = line[1:]
|
||||
if not line.startswith("!"):
|
||||
continue
|
||||
|
||||
pattern = line[1:].lstrip("/")
|
||||
while pattern.startswith("./"):
|
||||
pattern = pattern[2:]
|
||||
pattern = pattern.rstrip("/")
|
||||
parts = [part for part in pattern.split("/") if part and part != "."]
|
||||
if not parts:
|
||||
recurse_all = True
|
||||
continue
|
||||
|
||||
wildcard_index = next(
|
||||
(
|
||||
idx
|
||||
for idx, part in enumerate(parts)
|
||||
if any(char in part for char in _GLOB_CHARS)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if wildcard_index is not None:
|
||||
literal_parts = parts[:wildcard_index]
|
||||
if not literal_parts:
|
||||
recurse_all = True
|
||||
continue
|
||||
wildcard_prefixes.add(pathlib.PurePosixPath(*literal_parts))
|
||||
continue
|
||||
|
||||
parent_parts = parts[:-1]
|
||||
for idx in range(1, len(parent_parts) + 1):
|
||||
exact_dirs.add(pathlib.PurePosixPath(*parent_parts[:idx]))
|
||||
|
||||
return _NegatedDockerignoreHints(
|
||||
exact_dirs=frozenset(exact_dirs),
|
||||
wildcard_prefixes=frozenset(wildcard_prefixes),
|
||||
recurse_all=recurse_all,
|
||||
)
|
||||
@@ -9,12 +9,35 @@ from contextlib import contextmanager
|
||||
import click
|
||||
import pathspec
|
||||
|
||||
from langgraph_cli._ignore import _build_ignore_spec
|
||||
from langgraph_cli.config import Config, _assemble_local_deps
|
||||
|
||||
_WARN_SIZE = 50 * 1024 * 1024 # 50 MB
|
||||
_MAX_SIZE = 200 * 1024 * 1024 # 200 MB
|
||||
|
||||
_ALWAYS_EXCLUDE = [
|
||||
"__pycache__/",
|
||||
".git/",
|
||||
".venv/",
|
||||
"venv/",
|
||||
"node_modules/",
|
||||
".tox/",
|
||||
".mypy_cache/",
|
||||
]
|
||||
|
||||
|
||||
def _build_ignore_spec(directory: pathlib.Path) -> pathspec.PathSpec:
|
||||
"""Build a PathSpec combining built-in exclusions with .dockerignore and .gitignore.
|
||||
|
||||
Always excludes common non-source directories (_ALWAYS_EXCLUDE). On top of
|
||||
that, patterns from .dockerignore and .gitignore (if present) are merged in.
|
||||
"""
|
||||
lines: list[str] = list(_ALWAYS_EXCLUDE)
|
||||
for name in (".dockerignore", ".gitignore"):
|
||||
ignore_file = directory / name
|
||||
if ignore_file.is_file():
|
||||
lines.extend(ignore_file.read_text(encoding="utf-8").splitlines())
|
||||
return pathspec.PathSpec.from_lines("gitwildmatch", lines)
|
||||
|
||||
|
||||
def _tar_filter(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo | None:
|
||||
"""Strip symlinks, hardlinks, and traversal paths from archive."""
|
||||
|
||||
@@ -10,13 +10,7 @@ except ModuleNotFoundError: # pragma: no cover - exercised on Python 3.10.
|
||||
import tomli as tomllib
|
||||
|
||||
import click
|
||||
import pathspec
|
||||
|
||||
from langgraph_cli._ignore import (
|
||||
_build_dockerignore_negation_hints,
|
||||
_build_ignore_spec,
|
||||
_is_always_excluded,
|
||||
)
|
||||
from langgraph_cli.schemas import Config
|
||||
|
||||
|
||||
@@ -446,32 +440,16 @@ def _container_root_for_uv_lock_package(
|
||||
|
||||
|
||||
def _uv_lock_package_copy_items(
|
||||
package: UvLockPackage,
|
||||
plan: UvLockPlan,
|
||||
ignore_spec: pathspec.PathSpec,
|
||||
package: UvLockPackage, plan: UvLockPlan
|
||||
) -> tuple[tuple[pathlib.PurePosixPath, pathlib.PurePosixPath], ...]:
|
||||
# Skip entries that .dockerignore / built-in exclusions would strip from
|
||||
# the build context. Emitting `ADD <path>` for a file that Docker has
|
||||
# filtered out causes the build to fail with
|
||||
# "failed to compute cache key: <path> not found".
|
||||
if package.root != plan.project_root:
|
||||
relative_root = pathlib.PurePosixPath(
|
||||
*package.root.relative_to(plan.project_root).parts
|
||||
)
|
||||
if _is_always_excluded(relative_root, is_dir=True) or ignore_spec.match_file(
|
||||
f"{relative_root.as_posix()}/"
|
||||
):
|
||||
raise click.UsageError(
|
||||
f"Workspace member '{package.name}' at {relative_root} is "
|
||||
"excluded from the Docker build context, but uv.lock requires "
|
||||
"it to be copied into the build context. Remove the matching "
|
||||
"pattern or drop the member from [tool.uv.workspace].members."
|
||||
)
|
||||
return ((relative_root, plan.container_roots[package.root]),)
|
||||
|
||||
root_container = plan.container_roots[package.root]
|
||||
workspace_member_roots = plan.all_workspace_roots - {plan.project_root}
|
||||
negated_dockerignore_hints = _build_dockerignore_negation_hints(plan.project_root)
|
||||
|
||||
def iter_entries(
|
||||
current_dir: pathlib.Path,
|
||||
@@ -483,32 +461,18 @@ def _uv_lock_package_copy_items(
|
||||
# and excluded entirely otherwise.
|
||||
continue
|
||||
|
||||
descendant_member_roots = [
|
||||
ws_root
|
||||
for ws_root in workspace_member_roots
|
||||
if child in ws_root.parents
|
||||
]
|
||||
if child.is_dir() and descendant_member_roots:
|
||||
entries.extend(iter_entries(child))
|
||||
continue
|
||||
|
||||
relative_child = pathlib.PurePosixPath(
|
||||
*child.relative_to(plan.project_root).parts
|
||||
)
|
||||
is_dir = child.is_dir()
|
||||
if _is_always_excluded(relative_child, is_dir=is_dir):
|
||||
continue
|
||||
ignored = ignore_spec.match_file(
|
||||
f"{relative_child.as_posix()}/" if is_dir else relative_child.as_posix()
|
||||
)
|
||||
is_workspace_parent = is_dir and any(
|
||||
child in ws_root.parents for ws_root in workspace_member_roots
|
||||
)
|
||||
|
||||
if is_workspace_parent:
|
||||
entries.extend(iter_entries(child))
|
||||
continue
|
||||
if (
|
||||
is_dir
|
||||
and ignored
|
||||
and negated_dockerignore_hints.requires_dir_walk(relative_child)
|
||||
):
|
||||
entries.extend(iter_entries(child))
|
||||
continue
|
||||
if ignored:
|
||||
continue
|
||||
|
||||
entries.append(
|
||||
(relative_child, root_container.joinpath(*relative_child.parts))
|
||||
)
|
||||
@@ -992,13 +956,10 @@ def python_config_to_docker_uv_lock(
|
||||
docker_plan.add_raw("# -- End of uv.lock dependencies install --")
|
||||
docker_plan.add_blank()
|
||||
|
||||
ignore_spec = _build_ignore_spec(plan.project_root, include_gitignore=False)
|
||||
for package in plan.install_order:
|
||||
package_label = package.root.relative_to(plan.project_root).as_posix() or "."
|
||||
docker_plan.add_raw(f"# -- Adding workspace package {package_label} --")
|
||||
for source, destination in _uv_lock_package_copy_items(
|
||||
package, plan, ignore_spec
|
||||
):
|
||||
for source, destination in _uv_lock_package_copy_items(package, plan):
|
||||
docker_plan.add_raw(copy_from_project_root(source, destination.as_posix()))
|
||||
docker_plan.add_instruction(
|
||||
"WORKDIR", plan.container_roots[package.root].as_posix()
|
||||
|
||||
@@ -23,7 +23,7 @@ dependencies = [
|
||||
path = "langgraph_cli/__init__.py"
|
||||
[project.optional-dependencies]
|
||||
inmem = [
|
||||
"langgraph-api>=0.5.35,<0.9.0 ; python_version >= '3.11'",
|
||||
"langgraph-api>=0.5.35,<0.8.0 ; python_version >= '3.11'",
|
||||
"langgraph-runtime-inmem>=0.7 ; python_version >= '3.11'",
|
||||
]
|
||||
|
||||
|
||||
@@ -99,13 +99,6 @@ class TestBuildIgnoreSpec:
|
||||
assert spec.match_file("app.log")
|
||||
assert spec.match_file("mod.pyc")
|
||||
|
||||
def test_can_skip_gitignore(self, tmp_path):
|
||||
(tmp_path / ".dockerignore").write_text("*.log\n")
|
||||
(tmp_path / ".gitignore").write_text("*.pyc\n")
|
||||
spec = _build_ignore_spec(tmp_path, include_gitignore=False)
|
||||
assert spec.match_file("app.log")
|
||||
assert not spec.match_file("mod.pyc")
|
||||
|
||||
def test_no_ignore_files_only_builtins(self, tmp_path):
|
||||
spec = _build_ignore_spec(tmp_path)
|
||||
assert spec.match_file("__pycache__/")
|
||||
|
||||
@@ -4,7 +4,6 @@ import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
import textwrap
|
||||
from unittest.mock import patch
|
||||
|
||||
import click
|
||||
import pytest
|
||||
@@ -1856,364 +1855,6 @@ def test_config_to_docker_uv_lock_supports_single_uv_project_root():
|
||||
assert additional_contexts == {}
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_skips_dockerignore_entries():
|
||||
"""Entries filtered by .dockerignore / built-in excludes must not appear
|
||||
as ADD lines. Docker fails to compute the cache key for paths that the
|
||||
build context has stripped."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root = tmpdir_path / "single"
|
||||
project_root.mkdir()
|
||||
(project_root / "uv.lock").write_text("# uv lock file\n")
|
||||
(project_root / "pyproject.toml").write_text(
|
||||
textwrap.dedent(
|
||||
"""
|
||||
[project]
|
||||
name = "single-app"
|
||||
version = "0.1.0"
|
||||
dependencies = ["httpx>=0.28"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
"""
|
||||
).strip()
|
||||
+ "\n"
|
||||
)
|
||||
(project_root / "langgraph.json").write_text("{}\n")
|
||||
(project_root / "src").mkdir()
|
||||
(project_root / "src" / "agent.py").write_text("graph = object()\n")
|
||||
(project_root / "README.md").write_text("# hi\n")
|
||||
|
||||
# Built-in exclusions — must never appear as ADD lines.
|
||||
(project_root / ".git").mkdir()
|
||||
(project_root / ".git" / "HEAD").write_text("ref: refs/heads/main\n")
|
||||
(project_root / ".venv").mkdir()
|
||||
(project_root / ".venv" / "pyvenv.cfg").write_text("home = /usr\n")
|
||||
(project_root / "__pycache__").mkdir()
|
||||
(project_root / "__pycache__" / "x.cpython-311.pyc").write_bytes(b"\x00")
|
||||
|
||||
# .dockerignore excludes .gitignore and a custom path.
|
||||
(project_root / ".dockerignore").write_text(".gitignore\nsecrets.env\n")
|
||||
(project_root / ".gitignore").write_text("*.pyc\n")
|
||||
(project_root / "secrets.env").write_text("TOKEN=abc\n")
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {"agent": "./src/agent.py:graph"},
|
||||
"source": {"kind": "uv"},
|
||||
}
|
||||
)
|
||||
docker, _ = config_to_docker(
|
||||
project_root / "langgraph.json",
|
||||
config,
|
||||
base_image="langchain/langgraph-api:0.2.47",
|
||||
)
|
||||
|
||||
for excluded in (
|
||||
"ADD .git ",
|
||||
"ADD .gitignore ",
|
||||
"ADD .venv ",
|
||||
"ADD __pycache__ ",
|
||||
"ADD secrets.env ",
|
||||
):
|
||||
assert excluded not in docker, (
|
||||
f"{excluded!r} should be filtered out of Dockerfile:\n{docker}"
|
||||
)
|
||||
|
||||
# The .dockerignore itself is still part of the context and should be
|
||||
# ADDed (Docker needs it at build time, and archive.py includes it).
|
||||
assert "ADD .dockerignore /deps/workspace/.dockerignore" in docker
|
||||
assert "ADD src /deps/workspace/src" in docker
|
||||
assert "ADD README.md /deps/workspace/README.md" in docker
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_does_not_apply_gitignore():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root = tmpdir_path / "single"
|
||||
project_root.mkdir()
|
||||
(project_root / "uv.lock").write_text("# uv lock file\n")
|
||||
(project_root / "pyproject.toml").write_text(
|
||||
textwrap.dedent(
|
||||
"""
|
||||
[project]
|
||||
name = "single-app"
|
||||
version = "0.1.0"
|
||||
dependencies = ["httpx>=0.28"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
"""
|
||||
).strip()
|
||||
+ "\n"
|
||||
)
|
||||
(project_root / "langgraph.json").write_text("{}\n")
|
||||
(project_root / "src").mkdir()
|
||||
(project_root / "src" / "agent.py").write_text("graph = object()\n")
|
||||
(project_root / "README.md").write_text("# hi\n")
|
||||
(project_root / ".gitignore").write_text("README.md\n")
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {"agent": "./src/agent.py:graph"},
|
||||
"source": {"kind": "uv"},
|
||||
}
|
||||
)
|
||||
docker, _ = config_to_docker(
|
||||
project_root / "langgraph.json",
|
||||
config,
|
||||
base_image="langchain/langgraph-api:0.2.47",
|
||||
)
|
||||
|
||||
assert "ADD README.md /deps/workspace/README.md" in docker
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_skips_dockerignore_entries_in_workspace():
|
||||
"""Multi-member workspace: ignore patterns must filter root-level entries
|
||||
AND entries encountered while recursing into directories that contain
|
||||
workspace members (the `descendant_member_roots` branch)."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root, config_path = _write_uv_lock_workspace(
|
||||
tmpdir_path,
|
||||
agent_dependencies=["workspace-root", "shared", "httpx>=0.28"],
|
||||
root_sources="[tool.uv.sources]\nshared = { workspace = true }\nworkspace-root = { workspace = true }",
|
||||
agent_sources="[tool.uv.sources]\nshared = { workspace = true }\nworkspace-root = { workspace = true }",
|
||||
)
|
||||
root_src = project_root / "src" / "workspace_root"
|
||||
root_src.mkdir(parents=True)
|
||||
(root_src / "__init__.py").write_text("__all__ = []\n")
|
||||
(project_root / "README.md").write_text("workspace root package\n")
|
||||
|
||||
# A non-member sibling of the `apps/agent` member that should be
|
||||
# filtered out via .dockerignore. This exercises the recursion into
|
||||
# `apps/` where `apps/agent` is kept (it's a member) but its sibling is
|
||||
# filtered.
|
||||
(project_root / "apps" / "scratch.txt").write_text("scratch\n")
|
||||
# A root-level path that .dockerignore excludes.
|
||||
(project_root / "secrets.env").write_text("TOKEN=abc\n")
|
||||
(project_root / ".dockerignore").write_text("secrets.env\napps/scratch.txt\n")
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {
|
||||
"agent": "../../apps/agent/src/agent/graph.py:graph",
|
||||
},
|
||||
"source": {"kind": "uv", "root": "../..", "package": "agent"},
|
||||
}
|
||||
)
|
||||
docker, _ = config_to_docker(
|
||||
config_path, config, base_image="langchain/langgraph-api:0.2.47"
|
||||
)
|
||||
|
||||
assert "COPY --from=uv-workspace-root src /deps/workspace/src" in docker
|
||||
assert (
|
||||
"COPY --from=uv-workspace-root README.md /deps/workspace/README.md"
|
||||
in docker
|
||||
)
|
||||
assert (
|
||||
"COPY --from=uv-workspace-root .dockerignore /deps/workspace/.dockerignore"
|
||||
in docker
|
||||
)
|
||||
assert "secrets.env" not in docker
|
||||
assert "apps/scratch.txt" not in docker
|
||||
# Workspace members themselves are still copied via their own per-member
|
||||
# COPY line — the sibling filter must not disturb this.
|
||||
assert (
|
||||
"COPY --from=uv-workspace-root apps/agent /deps/workspace/apps/agent"
|
||||
in docker
|
||||
)
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_preserves_negated_dockerignore_descendants():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root = tmpdir_path / "single"
|
||||
project_root.mkdir()
|
||||
(project_root / "uv.lock").write_text("# uv lock file\n")
|
||||
(project_root / "pyproject.toml").write_text(
|
||||
textwrap.dedent(
|
||||
"""
|
||||
[project]
|
||||
name = "single-app"
|
||||
version = "0.1.0"
|
||||
dependencies = ["httpx>=0.28"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
"""
|
||||
).strip()
|
||||
+ "\n"
|
||||
)
|
||||
(project_root / "langgraph.json").write_text("{}\n")
|
||||
(project_root / "src").mkdir()
|
||||
(project_root / "src" / "agent.py").write_text("graph = object()\n")
|
||||
(project_root / "assets").mkdir()
|
||||
(project_root / "assets" / "keep.txt").write_text("keep\n")
|
||||
(project_root / "assets" / "drop.txt").write_text("drop\n")
|
||||
(project_root / ".dockerignore").write_text("assets/\n!assets/keep.txt\n")
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {"agent": "./src/agent.py:graph"},
|
||||
"source": {"kind": "uv"},
|
||||
}
|
||||
)
|
||||
docker, _ = config_to_docker(
|
||||
project_root / "langgraph.json",
|
||||
config,
|
||||
base_image="langchain/langgraph-api:0.2.47",
|
||||
)
|
||||
|
||||
assert "ADD assets /deps/workspace/assets" not in docker
|
||||
assert "ADD assets/keep.txt /deps/workspace/assets/keep.txt" in docker
|
||||
assert "assets/drop.txt" not in docker
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_prunes_unrelated_ignored_subtrees():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root = tmpdir_path / "single"
|
||||
project_root.mkdir()
|
||||
(project_root / "uv.lock").write_text("# uv lock file\n")
|
||||
(project_root / "pyproject.toml").write_text(
|
||||
textwrap.dedent(
|
||||
"""
|
||||
[project]
|
||||
name = "single-app"
|
||||
version = "0.1.0"
|
||||
dependencies = ["httpx>=0.28"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
"""
|
||||
).strip()
|
||||
+ "\n"
|
||||
)
|
||||
(project_root / "langgraph.json").write_text("{}\n")
|
||||
(project_root / "src").mkdir()
|
||||
(project_root / "src" / "agent.py").write_text("graph = object()\n")
|
||||
(project_root / "assets").mkdir()
|
||||
(project_root / "assets" / "keep.txt").write_text("keep\n")
|
||||
(project_root / "vendor").mkdir()
|
||||
(project_root / "vendor" / "huge.txt").write_text("large\n")
|
||||
(project_root / ".dockerignore").write_text(
|
||||
"vendor/\nassets/\n!assets/keep.txt\n"
|
||||
)
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {"agent": "./src/agent.py:graph"},
|
||||
"source": {"kind": "uv"},
|
||||
}
|
||||
)
|
||||
|
||||
original_iterdir = pathlib.Path.iterdir
|
||||
|
||||
def guarded_iterdir(self):
|
||||
if self == project_root / "vendor":
|
||||
raise AssertionError("should not walk unrelated ignored subtree")
|
||||
return original_iterdir(self)
|
||||
|
||||
with patch.object(
|
||||
pathlib.Path, "iterdir", autospec=True, side_effect=guarded_iterdir
|
||||
):
|
||||
docker, _ = config_to_docker(
|
||||
project_root / "langgraph.json",
|
||||
config,
|
||||
base_image="langchain/langgraph-api:0.2.47",
|
||||
)
|
||||
|
||||
assert "ADD assets/keep.txt /deps/workspace/assets/keep.txt" in docker
|
||||
assert "vendor/huge.txt" not in docker
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_never_reincludes_always_excluded_subtrees():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root = tmpdir_path / "single"
|
||||
project_root.mkdir()
|
||||
(project_root / "uv.lock").write_text("# uv lock file\n")
|
||||
(project_root / "pyproject.toml").write_text(
|
||||
textwrap.dedent(
|
||||
"""
|
||||
[project]
|
||||
name = "single-app"
|
||||
version = "0.1.0"
|
||||
dependencies = ["httpx>=0.28"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
"""
|
||||
).strip()
|
||||
+ "\n"
|
||||
)
|
||||
(project_root / "langgraph.json").write_text("{}\n")
|
||||
(project_root / "src").mkdir()
|
||||
(project_root / "src" / "agent.py").write_text("graph = object()\n")
|
||||
(project_root / ".venv" / "pkg").mkdir(parents=True)
|
||||
(project_root / ".venv" / "pkg" / "keep.txt").write_text("keep\n")
|
||||
(project_root / "node_modules" / "pkg").mkdir(parents=True)
|
||||
(project_root / "node_modules" / "pkg" / "package.json").write_text("{}\n")
|
||||
(project_root / ".dockerignore").write_text(
|
||||
"!.venv/pkg/keep.txt\n!node_modules/pkg/package.json\n"
|
||||
)
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {"agent": "./src/agent.py:graph"},
|
||||
"source": {"kind": "uv"},
|
||||
}
|
||||
)
|
||||
docker, _ = config_to_docker(
|
||||
project_root / "langgraph.json",
|
||||
config,
|
||||
base_image="langchain/langgraph-api:0.2.47",
|
||||
)
|
||||
|
||||
assert ".venv/pkg/keep.txt" not in docker
|
||||
assert "node_modules/pkg/package.json" not in docker
|
||||
assert "ADD src /deps/workspace/src" in docker
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_rejects_ignored_workspace_member():
|
||||
"""A workspace member matched by .dockerignore cannot be copied into the
|
||||
build context — uv.lock requires it, so fail loudly with a clear message."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
project_root, config_path = _write_uv_lock_workspace(
|
||||
tmpdir_path,
|
||||
agent_sources="[tool.uv.sources]\nshared = { workspace = true }",
|
||||
)
|
||||
(project_root / ".dockerignore").write_text("libs/shared\n")
|
||||
|
||||
config = validate_config(
|
||||
{
|
||||
"python_version": "3.11",
|
||||
"graphs": {"agent": "../../apps/agent/src/agent/graph.py:graph"},
|
||||
"source": {"kind": "uv", "root": "../..", "package": "agent"},
|
||||
"auth": {"path": "../../libs/shared/src/shared/auth.py:create_auth"},
|
||||
}
|
||||
)
|
||||
with pytest.raises(
|
||||
click.UsageError, match=r"Workspace member 'shared' at libs/shared"
|
||||
):
|
||||
config_to_docker(
|
||||
config_path, config, base_image="langchain/langgraph-api:0.2.47"
|
||||
)
|
||||
|
||||
|
||||
def test_config_to_docker_uv_lock_rejects_invalid_source_package_type():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_path = pathlib.Path(tmpdir)
|
||||
|
||||
Generated
+383
-465
File diff suppressed because it is too large
Load Diff
@@ -245,6 +245,15 @@ class _GraphCallbackManager(BaseCallbackManager):
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
def add_handler(
|
||||
self,
|
||||
handler: BaseCallbackHandler,
|
||||
inherit: bool = True, # noqa: FBT001,FBT002
|
||||
) -> None:
|
||||
if not isinstance(handler, GraphCallbackHandler):
|
||||
raise TypeError("handlers must inherit GraphCallbackHandler")
|
||||
super().add_handler(handler, inherit=inherit)
|
||||
|
||||
def copy(
|
||||
self,
|
||||
*,
|
||||
@@ -312,6 +321,15 @@ class _AsyncGraphCallbackManager(BaseCallbackManager):
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
def add_handler(
|
||||
self,
|
||||
handler: BaseCallbackHandler,
|
||||
inherit: bool = True, # noqa: FBT001,FBT002
|
||||
) -> None:
|
||||
if not isinstance(handler, GraphCallbackHandler):
|
||||
raise TypeError("handlers must inherit GraphCallbackHandler")
|
||||
super().add_handler(handler, inherit=inherit)
|
||||
|
||||
def copy(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from langgraph.channels.any_value import AnyValue
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
|
||||
from langgraph.channels.named_barrier_value import (
|
||||
@@ -20,6 +21,7 @@ __all__ = (
|
||||
"UntrackedValue",
|
||||
"EphemeralValue",
|
||||
"BinaryOperatorAggregate",
|
||||
"DeltaChannel",
|
||||
"NamedBarrierValue",
|
||||
"NamedBarrierValueAfterFinish",
|
||||
# topics
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Generic, TypeVar
|
||||
from typing import Any, Generic, Literal, TypeVar
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
@@ -119,3 +119,17 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
|
||||
Returns `True` if the channel was updated, `False` otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
|
||||
"""Called after checkpoint() with the assigned version, and after
|
||||
from_checkpoint() with the current channel version.
|
||||
|
||||
No-op by default. Override in channels that track their own version
|
||||
for incremental checkpointing (e.g. DeltaChannel).
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
def checkpoint_hydration_kind(self) -> Literal["delta"] | None:
|
||||
"""Return the saver hydration kind for this channel, if any."""
|
||||
return None
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import collections.abc
|
||||
from collections.abc import Callable, Sequence
|
||||
from copy import copy
|
||||
from typing import Any, Generic, Literal
|
||||
|
||||
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",)
|
||||
|
||||
|
||||
def _copy_value(value: Any) -> Any:
|
||||
if value is MISSING:
|
||||
return value
|
||||
try:
|
||||
return value.copy()
|
||||
except AttributeError:
|
||||
return copy(value)
|
||||
|
||||
|
||||
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]
|
||||
|
||||
@property
|
||||
def checkpoint_hydration_kind(self) -> Literal["delta"]:
|
||||
return "delta"
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
|
||||
new.key = self.key
|
||||
new.value = _copy_value(self.value)
|
||||
new._pending = self._pending[:]
|
||||
new._base_version = self._base_version
|
||||
new._last_checkpoint_id = self._last_checkpoint_id
|
||||
new._overwritten = self._overwritten
|
||||
new._steps_since_snapshot = self._steps_since_snapshot
|
||||
return new
|
||||
|
||||
def 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:
|
||||
pass
|
||||
elif isinstance(checkpoint, DeltaChainValue):
|
||||
accumulated: list[Value] = (
|
||||
checkpoint.base if checkpoint.base is not None else new.typ()
|
||||
)
|
||||
for step_writes in checkpoint.deltas:
|
||||
for write in step_writes:
|
||||
accumulated = new.operator(accumulated, write)
|
||||
new.value = accumulated
|
||||
# Seed the counter from actual chain depth so rehydration fires at
|
||||
# the right time regardless of how many prior invocations there were.
|
||||
new._steps_since_snapshot = len(checkpoint.deltas)
|
||||
elif isinstance(checkpoint, DeltaValue):
|
||||
# Should never reach here — checkpoint hydration should assemble
|
||||
# DeltaValues into DeltaChainValue before calling from_checkpoint.
|
||||
raise AssertionError(
|
||||
"DeltaChannel.from_checkpoint received a raw DeltaValue. "
|
||||
"This is a bug in checkpoint hydration — chain assembly should "
|
||||
"have occurred before from_checkpoint was called."
|
||||
)
|
||||
else:
|
||||
# Backwards compat: plain value from old BinaryOperatorAggregate checkpoint
|
||||
# or a full snapshot emitted by DeltaChannel.
|
||||
new.value = _copy_value(checkpoint)
|
||||
new._pending = []
|
||||
new._base_version = None # set by the subsequent after_checkpoint() call
|
||||
new._overwritten = False
|
||||
return new
|
||||
|
||||
def 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 = (
|
||||
_copy_value(overwrite_value)
|
||||
if overwrite_value is not None
|
||||
else self.typ()
|
||||
)
|
||||
self._pending = (
|
||||
[] if overwrite_value is None else [_copy_value(self.value)]
|
||||
)
|
||||
self._overwritten = True
|
||||
seen_overwrite = True
|
||||
elif not seen_overwrite:
|
||||
base = self.typ() if self.value is MISSING else self.value
|
||||
self.value = self.operator(base, value)
|
||||
self._pending.append(value)
|
||||
return True
|
||||
|
||||
def get(self) -> 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 _copy_value(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
|
||||
@@ -3,7 +3,11 @@ from __future__ import annotations
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from langgraph.checkpoint.base import Checkpoint
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointHydrationPlan,
|
||||
IncrementalChannelSpec,
|
||||
)
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
@@ -13,6 +17,19 @@ from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
LATEST_VERSION = 4
|
||||
|
||||
|
||||
def checkpoint_hydration_plan(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
) -> CheckpointHydrationPlan | None:
|
||||
"""Build a saver hydration plan from channel specs."""
|
||||
channels = tuple(
|
||||
IncrementalChannelSpec(name=name, kind=channel.checkpoint_hydration_kind)
|
||||
for name, channel in specs.items()
|
||||
if isinstance(channel, BaseChannel)
|
||||
and channel.checkpoint_hydration_kind is not None
|
||||
)
|
||||
return CheckpointHydrationPlan(channels=channels) if channels else None
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
v=LATEST_VERSION,
|
||||
@@ -67,13 +84,12 @@ def channels_from_checkpoint(
|
||||
channel_specs[k] = v
|
||||
else:
|
||||
managed_specs[k] = v
|
||||
return (
|
||||
{
|
||||
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
for k, v in channel_specs.items()
|
||||
},
|
||||
managed_specs,
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
|
||||
@@ -12,6 +12,7 @@ from contextlib import (
|
||||
ExitStack,
|
||||
)
|
||||
from datetime import datetime, timezone
|
||||
from functools import cached_property
|
||||
from inspect import signature
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
@@ -29,6 +30,7 @@ from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointHydrationPlan,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
PendingWrite,
|
||||
@@ -93,6 +95,7 @@ from langgraph.pregel._algo import (
|
||||
)
|
||||
from langgraph.pregel._checkpoint import (
|
||||
channels_from_checkpoint,
|
||||
checkpoint_hydration_plan,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
@@ -314,6 +317,29 @@ class PregelLoop:
|
||||
)
|
||||
self.prev_checkpoint_config = None
|
||||
|
||||
@cached_property
|
||||
def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None:
|
||||
"""Build the saver hydration plan from this loop's channel specs."""
|
||||
return checkpoint_hydration_plan(self.specs)
|
||||
|
||||
def _materialize_saved_checkpoint(
|
||||
self, saved: CheckpointTuple | None
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or self.checkpointer is None:
|
||||
return saved
|
||||
return self.checkpointer.materialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
async def _amaterialize_saved_checkpoint(
|
||||
self, saved: CheckpointTuple | None
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or self.checkpointer is None:
|
||||
return saved
|
||||
return await self.checkpointer.amaterialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
def _push_graph_lifecycle_event(
|
||||
self,
|
||||
kind: Literal["resume", "interrupt"],
|
||||
@@ -831,18 +857,8 @@ class PregelLoop:
|
||||
# parent. For forks (source=update/fork), use the fork's parent
|
||||
# checkpoint ID since the fork was created after the subgraph's
|
||||
# checkpoints from the original execution.
|
||||
#
|
||||
# Only gate on is_time_traveling (not is_replaying). When the
|
||||
# client resumes with an explicit checkpoint_id that happens to
|
||||
# point at the current head (e.g. LangGraph Studio sending
|
||||
# `checkpoint: {checkpoint_id}` alongside Command(resume=...)),
|
||||
# is_replaying is True but is_time_traveling is False. In that
|
||||
# case subgraphs should load their latest checkpoint normally,
|
||||
# not go through ReplayState's before-bound lookup which would
|
||||
# miss subgraph checkpoints created during processing of the
|
||||
# current parent step.
|
||||
replay_state: ReplayState | None = None
|
||||
if is_time_traveling:
|
||||
if self.is_replaying:
|
||||
replay_checkpoint_id = self.checkpoint["id"]
|
||||
if (
|
||||
self.checkpoint_metadata.get("source")
|
||||
@@ -891,6 +907,12 @@ 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()
|
||||
@@ -1247,6 +1269,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
|
||||
saved = self._materialize_saved_checkpoint(saved)
|
||||
|
||||
if saved is None:
|
||||
saved = CheckpointTuple(
|
||||
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
|
||||
@@ -1449,6 +1473,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
|
||||
saved = await self._amaterialize_saved_checkpoint(saved)
|
||||
|
||||
if saved is None:
|
||||
saved = CheckpointTuple(
|
||||
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
|
||||
|
||||
@@ -14,7 +14,7 @@ from langchain_core.messages import BaseMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph._internal._constants import NS_END, NS_SEP
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
from langgraph.types import Command
|
||||
@@ -132,15 +132,23 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
task_checkpoint_ns = cast(str, metadata["langgraph_checkpoint_ns"])
|
||||
checkpoint_ns = (
|
||||
f"{task_checkpoint_ns.rsplit(NS_END, 1)[0]}{NS_END}"
|
||||
if NS_END in task_checkpoint_ns
|
||||
else task_checkpoint_ns
|
||||
)
|
||||
ns = tuple(task_checkpoint_ns.split(NS_SEP))[:-1]
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
stream_metadata = dict(metadata)
|
||||
stream_metadata["langgraph_checkpoint_ns"] = checkpoint_ns
|
||||
# Preserve backwards-compatible streamed checkpoint metadata shape.
|
||||
stream_metadata["checkpoint_ns"] = checkpoint_ns
|
||||
if tags:
|
||||
if filtered_tags := [t for t in tags if not t.startswith("seq:step")]:
|
||||
metadata["tags"] = filtered_tags
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
stream_metadata["tags"] = filtered_tags
|
||||
self.metadata[run_id] = (ns, stream_metadata)
|
||||
|
||||
def on_llm_new_token(
|
||||
self,
|
||||
|
||||
@@ -17,7 +17,7 @@ from collections.abc import (
|
||||
Sequence,
|
||||
)
|
||||
from dataclasses import is_dataclass, replace
|
||||
from functools import partial
|
||||
from functools import cached_property, partial
|
||||
from inspect import isclass
|
||||
from typing import (
|
||||
Any,
|
||||
@@ -44,6 +44,7 @@ from langgraph.cache.base import BaseCache
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
CheckpointHydrationPlan,
|
||||
CheckpointTuple,
|
||||
)
|
||||
from langgraph.store.base import BaseStore
|
||||
@@ -123,6 +124,7 @@ from langgraph.pregel._algo import (
|
||||
from langgraph.pregel._call import identifier
|
||||
from langgraph.pregel._checkpoint import (
|
||||
channels_from_checkpoint,
|
||||
checkpoint_hydration_plan,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
@@ -728,6 +730,54 @@ class Pregel(
|
||||
return checkpointer
|
||||
return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist)
|
||||
|
||||
@cached_property
|
||||
def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None:
|
||||
return checkpoint_hydration_plan(self.channels)
|
||||
|
||||
def _materialize_saved_checkpoint(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: CheckpointTuple | None,
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or checkpointer is None:
|
||||
return saved
|
||||
return checkpointer.materialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
async def _amaterialize_saved_checkpoint(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: CheckpointTuple | None,
|
||||
) -> CheckpointTuple | None:
|
||||
if saved is None or checkpointer is None:
|
||||
return saved
|
||||
return await checkpointer.amaterialize_checkpoint_tuple(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
def _materialize_saved_checkpoints(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: Sequence[CheckpointTuple],
|
||||
) -> list[CheckpointTuple]:
|
||||
if checkpointer is None or not saved:
|
||||
return list(saved)
|
||||
return checkpointer.materialize_checkpoint_tuples(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
async def _amaterialize_saved_checkpoints(
|
||||
self,
|
||||
checkpointer: BaseCheckpointSaver | None,
|
||||
saved: Sequence[CheckpointTuple],
|
||||
) -> list[CheckpointTuple]:
|
||||
if checkpointer is None or not saved:
|
||||
return list(saved)
|
||||
return await checkpointer.amaterialize_checkpoint_tuples(
|
||||
saved, self._checkpoint_hydration_plan
|
||||
)
|
||||
|
||||
def get_graph(
|
||||
self, config: RunnableConfig | None = None, *, xray: int | bool = False
|
||||
) -> Graph:
|
||||
@@ -1049,13 +1099,14 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1168,13 +1219,14 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1300,7 +1352,9 @@ class Pregel(
|
||||
if not isinstance(thread_id, str):
|
||||
config[CONF][CONFIG_KEY_THREAD_ID] = str(thread_id)
|
||||
|
||||
saved = checkpointer.get_tuple(config)
|
||||
saved = self._materialize_saved_checkpoint(
|
||||
checkpointer, checkpointer.get_tuple(config)
|
||||
)
|
||||
return self._prepare_state_snapshot(
|
||||
config,
|
||||
saved,
|
||||
@@ -1344,7 +1398,9 @@ class Pregel(
|
||||
if not isinstance(thread_id, str):
|
||||
config[CONF][CONFIG_KEY_THREAD_ID] = str(thread_id)
|
||||
|
||||
saved = await checkpointer.aget_tuple(config)
|
||||
saved = await self._amaterialize_saved_checkpoint(
|
||||
checkpointer, await checkpointer.aget_tuple(config)
|
||||
)
|
||||
return await self._aprepare_state_snapshot(
|
||||
config,
|
||||
saved,
|
||||
@@ -1398,9 +1454,11 @@ class Pregel(
|
||||
},
|
||||
)
|
||||
# eagerly consume list() to avoid holding up the db cursor
|
||||
for checkpoint_tuple in list(
|
||||
checkpointer.list(config, before=before, limit=limit, filter=filter)
|
||||
):
|
||||
checkpoint_tuples = self._materialize_saved_checkpoints(
|
||||
checkpointer,
|
||||
list(checkpointer.list(config, before=before, limit=limit, filter=filter)),
|
||||
)
|
||||
for checkpoint_tuple in checkpoint_tuples:
|
||||
yield self._prepare_state_snapshot(
|
||||
checkpoint_tuple.config, checkpoint_tuple
|
||||
)
|
||||
@@ -1452,12 +1510,16 @@ class Pregel(
|
||||
},
|
||||
)
|
||||
# eagerly consume list() to avoid holding up the db cursor
|
||||
for checkpoint_tuple in [
|
||||
c
|
||||
async for c in checkpointer.alist(
|
||||
config, before=before, limit=limit, filter=filter
|
||||
)
|
||||
]:
|
||||
checkpoint_tuples = await self._amaterialize_saved_checkpoints(
|
||||
checkpointer,
|
||||
[
|
||||
c
|
||||
async for c in checkpointer.alist(
|
||||
config, before=before, limit=limit, filter=filter
|
||||
)
|
||||
],
|
||||
)
|
||||
for checkpoint_tuple in checkpoint_tuples:
|
||||
yield await self._aprepare_state_snapshot(
|
||||
checkpoint_tuple.config, checkpoint_tuple
|
||||
)
|
||||
@@ -1517,12 +1579,13 @@ class Pregel(
|
||||
) -> RunnableConfig:
|
||||
# get last checkpoint
|
||||
config = ensure_config(self.config, input_config)
|
||||
saved = checkpointer.get_tuple(config)
|
||||
saved = self._materialize_saved_checkpoint(
|
||||
checkpointer, checkpointer.get_tuple(config)
|
||||
)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
@@ -1963,12 +2026,13 @@ class Pregel(
|
||||
) -> RunnableConfig:
|
||||
# get last checkpoint
|
||||
config = ensure_config(self.config, input_config)
|
||||
saved = await checkpointer.aget_tuple(config)
|
||||
saved = await self._amaterialize_saved_checkpoint(
|
||||
checkpointer, await checkpointer.aget_tuple(config)
|
||||
)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.1.7a2"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
@@ -24,10 +24,10 @@ classifiers = [
|
||||
'Programming Language :: Python :: 3.13',
|
||||
]
|
||||
dependencies = [
|
||||
"langchain-core>=1.3.0,<2",
|
||||
"langchain-core==1.3.0a2",
|
||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||
"langgraph-sdk>=0.3.0,<0.4.0",
|
||||
"langgraph-prebuilt>=1.0.12,<1.1.0",
|
||||
"langgraph-prebuilt>=1.0.9,<1.1.0",
|
||||
"xxhash>=3.5.0",
|
||||
"pydantic>=2.7.4",
|
||||
]
|
||||
|
||||
@@ -117,3 +117,609 @@ 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 DeltaChainValue, 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 len(d.delta[0]) == 1
|
||||
assert d.delta[0][0].content == "new"
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
replayed = spec.from_checkpoint(DeltaChainValue(base=None, deltas=[d.delta]))
|
||||
assert replayed.get()[0].content == "new"
|
||||
|
||||
|
||||
def test_delta_channel_assembly_fallback_via_get_tuple() -> None:
|
||||
"""Materialization falls back to get_tuple for savers without get_channel_blob."""
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
CheckpointHydrationPlan,
|
||||
CheckpointTuple,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
IncrementalChannelSpec,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
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"
|
||||
)
|
||||
|
||||
class TestSaver(BaseCheckpointSaver[str]):
|
||||
def get_tuple(self, config):
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "t1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "cp1",
|
||||
}
|
||||
},
|
||||
checkpoint=cp1,
|
||||
metadata={},
|
||||
parent_config=None,
|
||||
pending_writes=[],
|
||||
)
|
||||
|
||||
def list(self, config, *, filter=None, before=None, limit=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def put(self, config, checkpoint, metadata, new_versions):
|
||||
raise NotImplementedError
|
||||
|
||||
saver = TestSaver()
|
||||
|
||||
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
|
||||
materialized = saver.materialize_checkpoint(
|
||||
config,
|
||||
cp2,
|
||||
CheckpointHydrationPlan(
|
||||
channels=(IncrementalChannelSpec(name="messages", kind="delta"),)
|
||||
),
|
||||
)
|
||||
|
||||
chain = materialized["channel_values"]["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_dict_reducer_overwrite_preserves_mapping() -> None:
|
||||
"""Overwrite should preserve dict values instead of coercing them to keys."""
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
ch = DeltaChannel(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
ch.update([{"a": 1}])
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
ch.update([Overwrite({"b": 2})])
|
||||
d = ch.checkpoint()
|
||||
assert isinstance(d, DeltaValue)
|
||||
assert d.prev_checkpoint_id is None
|
||||
assert d.delta == [{"b": 2}]
|
||||
assert ch.get() == {"b": 2}
|
||||
|
||||
spec = DeltaChannel(merge_dicts, dict)
|
||||
replayed = spec.from_checkpoint(DeltaChainValue(base=None, deltas=[d.delta]))
|
||||
assert replayed.get() == {"b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_dict_snapshot_every_round_trip() -> None:
|
||||
"""Full snapshots should preserve non-list reducers across reload."""
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
ch = DeltaChannel(merge_dicts, dict, snapshot_every=1).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint("v0", checkpoint_id="cid0")
|
||||
|
||||
ch.update([{"a": 1}])
|
||||
first = ch.checkpoint()
|
||||
assert isinstance(first, DeltaValue)
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
ch.update([{"b": 2}])
|
||||
snap = ch.checkpoint()
|
||||
assert isinstance(snap, dict)
|
||||
assert snap == {"a": 1, "b": 2}
|
||||
|
||||
rehydrated = DeltaChannel(merge_dicts, dict, snapshot_every=1).from_checkpoint(snap)
|
||||
assert rehydrated.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_assembly_fast_path_returns_delta_value() -> None:
|
||||
"""get_channel_blob returning a DeltaValue continues chain traversal (fast-path)."""
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
CheckpointHydrationPlan,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
IncrementalChannelSpec,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
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"
|
||||
)
|
||||
|
||||
class TestSaver(BaseCheckpointSaver[str]):
|
||||
def get_tuple(self, config):
|
||||
raise NotImplementedError
|
||||
|
||||
def list(self, config, *, filter=None, before=None, limit=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def put(self, config, checkpoint, metadata, new_versions):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_channel_blob(self, thread_id, checkpoint_ns, checkpoint_id, channel):
|
||||
if checkpoint_id == "cp2":
|
||||
return dv_cp2
|
||||
if checkpoint_id == "cp1":
|
||||
return [msg1]
|
||||
return NotImplemented
|
||||
|
||||
saver = TestSaver()
|
||||
|
||||
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
|
||||
materialized = saver.materialize_checkpoint(
|
||||
config,
|
||||
cp3,
|
||||
CheckpointHydrationPlan(
|
||||
channels=(IncrementalChannelSpec(name="messages", kind="delta"),)
|
||||
),
|
||||
)
|
||||
|
||||
chain = materialized["channel_values"]["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_dict_reducer_fresh_channel() -> None:
|
||||
"""DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint."""
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
ch = DeltaChannel(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
# Should be available (not raise EmptyChannelError) and start empty
|
||||
assert ch.is_available()
|
||||
assert ch.get() == {}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_basic_updates() -> None:
|
||||
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
ch = DeltaChannel(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
ch.update([{"a": 1}])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
assert d1.delta == [{"a": 1}]
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
ch.update([{"b": 2}])
|
||||
d2 = ch.checkpoint()
|
||||
assert d2.delta == [{"b": 2}]
|
||||
assert d2.prev_checkpoint_id == "cid1"
|
||||
ch.after_checkpoint("v2")
|
||||
|
||||
assert ch.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_chain_reconstruction() -> None:
|
||||
"""DeltaChainValue replays correctly through a dict merge reducer."""
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
spec = DeltaChannel(merge_dicts, dict)
|
||||
chain = DeltaChainValue(
|
||||
base={"a": 1},
|
||||
deltas=[[{"b": 2}], [{"c": 3}]],
|
||||
)
|
||||
ch = spec.from_checkpoint(chain)
|
||||
assert ch.get() == {"a": 1, "b": 2, "c": 3}
|
||||
assert ch._steps_since_snapshot == 2
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_deletions() -> None:
|
||||
"""Dict reducer that treats None values as deletions works end-to-end (deepagents pattern)."""
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
|
||||
def merge_files(left: dict | None, right: dict) -> dict:
|
||||
if left is None:
|
||||
return {k: v for k, v in right.items() if v is not None}
|
||||
result = {**left}
|
||||
for k, v in right.items():
|
||||
if v is None:
|
||||
result.pop(k, None)
|
||||
else:
|
||||
result[k] = v
|
||||
return result
|
||||
|
||||
ch = DeltaChannel(merge_files, dict).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
# Delete file1, add file3
|
||||
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
||||
ch.after_checkpoint("v2", checkpoint_id="cid2")
|
||||
|
||||
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
# Confirm chain reconstruction produces the same result
|
||||
chain = DeltaChainValue(
|
||||
base={},
|
||||
deltas=[
|
||||
[{"file1.py": "content1", "file2.py": "content2"}],
|
||||
[{"file1.py": None, "file3.py": "content3"}],
|
||||
],
|
||||
)
|
||||
spec = DeltaChannel(merge_files, dict)
|
||||
ch2 = spec.from_checkpoint(chain)
|
||||
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
|
||||
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 langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
CheckpointHydrationPlan,
|
||||
DeltaValue,
|
||||
IncrementalChannelSpec,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
cp = empty_checkpoint()
|
||||
cp["id"] = "cp2"
|
||||
cp["channel_values"]["messages"] = DeltaValue(
|
||||
delta=["msg2"], prev_checkpoint_id="cp-missing"
|
||||
)
|
||||
|
||||
class TestSaver(BaseCheckpointSaver[str]):
|
||||
def get_tuple(self, config):
|
||||
return None
|
||||
|
||||
def list(self, config, *, filter=None, before=None, limit=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def put(self, config, checkpoint, metadata, new_versions):
|
||||
raise NotImplementedError
|
||||
|
||||
saver = TestSaver()
|
||||
|
||||
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
|
||||
|
||||
materialized = saver.materialize_checkpoint(
|
||||
config,
|
||||
cp,
|
||||
CheckpointHydrationPlan(
|
||||
channels=(IncrementalChannelSpec(name="messages", kind="delta"),)
|
||||
),
|
||||
)
|
||||
|
||||
# Should still assemble — with partial chain (just the current delta, base=None)
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
|
||||
chain = materialized["channel_values"]["messages"]
|
||||
assert isinstance(chain, DeltaChainValue)
|
||||
assert chain.base is None
|
||||
assert chain.deltas == [["msg2"]]
|
||||
|
||||
@@ -0,0 +1,356 @@
|
||||
"""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,70 +275,3 @@ def test_graph_callbacks_accept_base_callback_manager() -> None:
|
||||
|
||||
assert "__interrupt__" in first
|
||||
assert len(graph_handler.interrupt_events) == 1
|
||||
|
||||
|
||||
def test_non_graph_handler_via_add_handler_does_not_crash() -> None:
|
||||
"""Non-GraphCallbackHandler added via add_handler should not raise.
|
||||
|
||||
Libraries like opentelemetry-instrumentation-langchain monkey-patch
|
||||
BaseCallbackManager.__init__ and inject handlers via add_handler().
|
||||
These handlers inherit from BaseCallbackHandler, not
|
||||
GraphCallbackHandler. They must be silently accepted — graph lifecycle
|
||||
events will simply not be dispatched to them.
|
||||
"""
|
||||
from langgraph.callbacks import _GraphCallbackManager
|
||||
|
||||
manager = _GraphCallbackManager()
|
||||
plain_handler = _LangChainCustomEventHandler()
|
||||
|
||||
manager.add_handler(plain_handler, inherit=True)
|
||||
assert plain_handler in manager.handlers
|
||||
|
||||
|
||||
def test_non_graph_handler_does_not_receive_lifecycle_events() -> None:
|
||||
"""Non-GraphCallbackHandler added alongside a GraphCallbackHandler
|
||||
should not interfere with lifecycle event dispatch."""
|
||||
graph = _build_interrupt_graph()
|
||||
graph_handler = _GraphEventHandler()
|
||||
plain_handler = _LangChainCustomEventHandler()
|
||||
|
||||
config = {
|
||||
"configurable": {"thread_id": "graph-callback-mixed-handlers"},
|
||||
"callbacks": [plain_handler, graph_handler],
|
||||
}
|
||||
|
||||
first = graph.invoke({"answer": None}, config)
|
||||
assert "__interrupt__" in first
|
||||
|
||||
assert len(graph_handler.interrupt_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
resumed = graph.invoke(Command(resume="done"), config)
|
||||
assert resumed == {"answer": "done"}
|
||||
assert len(graph_handler.resume_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_non_graph_handler_does_not_receive_lifecycle_events_async() -> None:
|
||||
"""Async variant: non-GraphCallbackHandler should not interfere."""
|
||||
graph = _build_interrupt_graph()
|
||||
graph_handler = _GraphEventHandler()
|
||||
plain_handler = _LangChainCustomEventHandler()
|
||||
|
||||
config = {
|
||||
"configurable": {"thread_id": "graph-callback-mixed-handlers-async"},
|
||||
"callbacks": [plain_handler, graph_handler],
|
||||
}
|
||||
|
||||
first = await graph.ainvoke({"answer": None}, config)
|
||||
assert "__interrupt__" in first
|
||||
|
||||
assert len(graph_handler.interrupt_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
resumed = await graph.ainvoke(Command(resume="done"), config)
|
||||
assert resumed == {"answer": "done"}
|
||||
assert len(graph_handler.resume_events) == 1
|
||||
assert plain_handler.events == []
|
||||
|
||||
@@ -9400,3 +9400,190 @@ 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"
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
"""Sweep snapshot_every values to find the storage vs. time-travel tradeoff.
|
||||
|
||||
Run directly: python tests/test_rehydrate_sweep.py
|
||||
Run via pytest: pytest tests/test_rehydrate_sweep.py -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
REHYDRATE_SWEEP = [5, 10, 25, 50, 100, None] # None = no rehydration (pure diff)
|
||||
TURN_COUNTS = [50, 100, 250, 500]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_state(snapshot_every: int | None) -> type:
|
||||
channel = DeltaChannel(add_messages, snapshot_every=snapshot_every)
|
||||
return TypedDict("S", {"messages": Annotated[list, channel]})
|
||||
|
||||
|
||||
def _make_graph(state_cls: type) -> Any:
|
||||
def human_node(state: Any) -> dict:
|
||||
return {}
|
||||
|
||||
def ai_node(state: Any) -> dict:
|
||||
last = state["messages"][-1]
|
||||
return {"messages": [AIMessage(content=f"reply-to-{last.id}")]}
|
||||
|
||||
g = StateGraph(state_cls)
|
||||
g.add_node("human", human_node)
|
||||
g.add_node("ai", ai_node)
|
||||
g.add_edge("human", "ai")
|
||||
g.add_edge("ai", END)
|
||||
g.set_entry_point("human")
|
||||
return g.compile(checkpointer=MemorySaver())
|
||||
|
||||
|
||||
def _total_blob_bytes(saver: MemorySaver) -> int:
|
||||
total = 0
|
||||
for (_, _, _, _), (type_tag, blob) in saver.blobs.items():
|
||||
if blob is not None:
|
||||
total += len(blob)
|
||||
return total
|
||||
|
||||
|
||||
def _measure_time_travel_ms(graph: Any, config: dict) -> float:
|
||||
"""Time how long it takes to get state at the very first checkpoint (worst case)."""
|
||||
history = list(graph.get_state_history(config))
|
||||
if not history:
|
||||
return 0.0
|
||||
oldest = history[-1]
|
||||
t0 = time.perf_counter()
|
||||
graph.get_state(oldest.config)
|
||||
return (time.perf_counter() - t0) * 1000
|
||||
|
||||
|
||||
def _run(n_turns: int, snapshot_every: int | None) -> tuple[float, int, float]:
|
||||
"""Returns (write_ms, blob_bytes, time_travel_ms)."""
|
||||
state_cls = _make_state(snapshot_every)
|
||||
graph = _make_graph(state_cls)
|
||||
saver: MemorySaver = graph.checkpointer # type: ignore[assignment]
|
||||
config = {"configurable": {"thread_id": "sweep"}}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
for i in range(n_turns):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"msg-{i}", id=f"h{i}")]}, config
|
||||
)
|
||||
write_ms = (time.perf_counter() - t0) * 1000
|
||||
|
||||
blob_bytes = _total_blob_bytes(saver)
|
||||
tt_ms = _measure_time_travel_ms(graph, config)
|
||||
return write_ms, blob_bytes, tt_ms
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ASCII sparkline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sparkline(values: list[float], width: int = 20) -> str:
|
||||
bars = " ▁▂▃▄▅▆▇█"
|
||||
lo, hi = min(values), max(values)
|
||||
span = hi - lo or 1
|
||||
chars = [bars[round((v - lo) / span * (len(bars) - 1))] for v in values]
|
||||
return "".join(chars).ljust(width)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_sweep() -> None:
|
||||
label = {v: (str(v) if v is not None else "None(∞)") for v in REHYDRATE_SWEEP}
|
||||
|
||||
print()
|
||||
print("snapshot_every sweep — storage vs time-travel cost")
|
||||
print("=" * 90)
|
||||
|
||||
for turns in TURN_COUNTS:
|
||||
print(f"\n--- {turns} turns ---")
|
||||
col_w = 12
|
||||
header = (
|
||||
f"{'snapshot_every':>18} "
|
||||
f"{'blob_bytes':>{col_w}} "
|
||||
f"{'write_ms':>{col_w}} "
|
||||
f"{'time_travel_ms':>{col_w}}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * 60)
|
||||
|
||||
tt_vals: list[float] = []
|
||||
byte_vals: list[int] = []
|
||||
write_vals: list[float] = []
|
||||
rows: list[tuple] = []
|
||||
|
||||
for rv in REHYDRATE_SWEEP:
|
||||
write_ms, blob_bytes, tt_ms = _run(turns, rv)
|
||||
rows.append((rv, blob_bytes, write_ms, tt_ms))
|
||||
byte_vals.append(blob_bytes)
|
||||
write_vals.append(write_ms)
|
||||
tt_vals.append(tt_ms)
|
||||
|
||||
for rv, blob_bytes, write_ms, tt_ms in rows:
|
||||
print(
|
||||
f"{label[rv]:>18} "
|
||||
f"{blob_bytes:>{col_w},} "
|
||||
f"{write_ms:>{col_w}.1f} "
|
||||
f"{tt_ms:>{col_w}.2f}"
|
||||
)
|
||||
|
||||
print()
|
||||
print(
|
||||
f" bytes spark: [{_sparkline(byte_vals)}] "
|
||||
f"lo={min(byte_vals):,} hi={max(byte_vals):,}"
|
||||
)
|
||||
print(
|
||||
f" time-travel spark: [{_sparkline(tt_vals)}] "
|
||||
f"lo={min(tt_vals):.2f}ms hi={max(tt_vals):.2f}ms"
|
||||
)
|
||||
print(
|
||||
f" write spark: [{_sparkline(write_vals)}] "
|
||||
f"lo={min(write_vals):.1f}ms hi={max(write_vals):.1f}ms"
|
||||
)
|
||||
|
||||
print()
|
||||
print("=" * 90)
|
||||
print(
|
||||
"snapshot_every=None means pure diff (no snapshots) — "
|
||||
"lowest storage, highest time-travel cost."
|
||||
)
|
||||
print(
|
||||
"Lower snapshot_every = more frequent full snapshots = "
|
||||
"faster time-travel, more storage."
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
def test_rehydrate_sweep(capsys: Any) -> None:
|
||||
with capsys.disabled():
|
||||
run_sweep()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_sweep()
|
||||
sys.exit(0)
|
||||
@@ -1113,70 +1113,6 @@ def test_subgraph_interrupt_replay_from_parent_then_resume(
|
||||
]
|
||||
|
||||
|
||||
def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Resume with Command(resume=...) plus the current head checkpoint_id
|
||||
in config. The subgraph must continue from the interrupted node, not
|
||||
restart from scratch. Explicit checkpoint_id triggers is_replaying but
|
||||
this is a resume, not a time-travel, so ReplayState should not apply."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["sub_a"]}
|
||||
|
||||
def ask_human(state: State) -> State:
|
||||
called.append("ask_human")
|
||||
answer = interrupt("Provide input:")
|
||||
return {"value": [f"human:{answer}"]}
|
||||
|
||||
def step_b(state: State) -> State:
|
||||
called.append("step_b")
|
||||
return {"value": ["sub_b"]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_human", ask_human)
|
||||
.add_node("step_b", step_b)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_human")
|
||||
.add_edge("ask_human", "step_b")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("subgraph_node", subgraph)
|
||||
.add_edge(START, "subgraph_node")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until interrupt fires in subgraph
|
||||
graph.invoke({"value": []}, config)
|
||||
assert called == ["step_a", "ask_human"]
|
||||
|
||||
# Resume with explicit head checkpoint_id in config
|
||||
head_checkpoint_id = graph.get_state(config).config["configurable"]["checkpoint_id"]
|
||||
called.clear()
|
||||
resume_config = {
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_id": head_checkpoint_id,
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
result = graph.invoke(Command(resume="answer"), resume_config)
|
||||
|
||||
assert called == ["ask_human", "step_b"]
|
||||
assert "__interrupt__" not in result
|
||||
assert result["value"] == ["sub_a", "human:answer", "sub_b"]
|
||||
|
||||
|
||||
def test_subgraph_replay_loads_accumulated_state_then_resume(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
|
||||
Generated
+15
-15
@@ -1348,7 +1348,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.0a2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -1360,14 +1360,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.1.7a2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1439,7 +1439,7 @@ test = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
@@ -1548,7 +1548,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1706,7 +1706,7 @@ inmem = [
|
||||
requires-dist = [
|
||||
{ name = "click", specifier = ">=8.1.7" },
|
||||
{ name = "httpx", specifier = ">=0.24.0" },
|
||||
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.9.0" },
|
||||
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.8.0" },
|
||||
{ name = "langgraph-runtime-inmem", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.7" },
|
||||
{ name = "langgraph-sdk", marker = "python_full_version >= '3.11'", specifier = ">=0.1.0" },
|
||||
{ name = "pathspec", specifier = ">=0.11.0" },
|
||||
@@ -1742,7 +1742,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.12"
|
||||
version = "1.0.9"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1751,7 +1751,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.1" },
|
||||
{ name = "langchain-core", specifier = ">=1.0.0" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
]
|
||||
|
||||
@@ -2140,7 +2140,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "nbconvert"
|
||||
version = "7.17.1"
|
||||
version = "7.17.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "beautifulsoup4" },
|
||||
@@ -2158,9 +2158,9 @@ dependencies = [
|
||||
{ name = "pygments" },
|
||||
{ name = "traitlets" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/01/b1/708e53fe2e429c103c6e6e159106bcf0357ac41aa4c28772bd8402339051/nbconvert-7.17.1.tar.gz", hash = "sha256:34d0d0a7e73ce3cbab6c5aae8f4f468797280b01fd8bd2ca746da8569eddd7d2", size = 865311, upload-time = "2026-04-08T00:44:14.914Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/38/47/81f886b699450d0569f7bc551df2b1673d18df7ff25cc0c21ca36ed8a5ff/nbconvert-7.17.0.tar.gz", hash = "sha256:1b2696f1b5be12309f6c7d707c24af604b87dfaf6d950794c7b07acab96dda78", size = 862855, upload-time = "2026-01-29T16:37:48.478Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/67/f8/bb0a9d5f46819c821dc1f004aa2cc29b1d91453297dbf5ff20470f00f193/nbconvert-7.17.1-py3-none-any.whl", hash = "sha256:aa85c087b435e7bf1ffd03319f658e285f2b89eccab33bc1ba7025495ab3e7c8", size = 261927, upload-time = "2026-04-08T00:44:12.845Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0d/4b/8d5f796a792f8a25f6925a96032f098789f448571eb92011df1ae59e8ea8/nbconvert-7.17.0-py3-none-any.whl", hash = "sha256:4f99a63b337b9a23504347afdab24a11faa7d86b405e5c8f9881cd313336d518", size = 261510, upload-time = "2026-01-29T16:37:46.322Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3018,11 +3018,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "python-dotenv"
|
||||
version = "1.2.2"
|
||||
version = "1.2.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/82/ed/0301aeeac3e5353ef3d94b6ec08bbcabd04a72018415dcb29e588514bba8/python_dotenv-1.2.2.tar.gz", hash = "sha256:2c371a91fbd7ba082c2c1dc1f8bf89ca22564a087c2c287cd9b662adde799cf3", size = 50135, upload-time = "2026-03-01T16:00:26.196Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f0/26/19cadc79a718c5edbec86fd4919a6b6d3f681039a2f6d66d14be94e75fb9/python_dotenv-1.2.1.tar.gz", hash = "sha256:42667e897e16ab0d66954af0e60a9caa94f0fd4ecf3aaf6d2d260eec1aa36ad6", size = 44221, upload-time = "2025-10-26T15:12:10.434Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101, upload-time = "2026-03-01T16:00:25.09Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/14/1b/a298b06749107c305e1fe0f814c6c74aea7b2f1e10989cb30f544a1b3253/python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61", size = 21230, upload-time = "2025-10-26T15:12:09.109Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -82,7 +82,6 @@ from langchain_core.tools.base import (
|
||||
_is_injected_arg_type,
|
||||
get_all_basemodel_annotations,
|
||||
)
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_READ
|
||||
from langgraph._internal._runnable import RunnableCallable
|
||||
from langgraph.errors import GraphBubbleUp
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES
|
||||
@@ -801,7 +800,7 @@ class ToolNode(RunnableCallable):
|
||||
# Construct ToolRuntime instances at the top level for each tool call
|
||||
tool_runtimes = []
|
||||
for call, cfg in zip(tool_calls, config_list, strict=False):
|
||||
state = self._extract_state(input, cfg)
|
||||
state = self._extract_state(input)
|
||||
tool_runtime = ToolRuntime(
|
||||
state=state,
|
||||
tool_call_id=call["id"],
|
||||
@@ -809,7 +808,6 @@ class ToolNode(RunnableCallable):
|
||||
context=runtime.context,
|
||||
store=runtime.store,
|
||||
stream_writer=runtime.stream_writer,
|
||||
tools=list(self.tools_by_name.values()),
|
||||
execution_info=runtime.execution_info,
|
||||
server_info=runtime.server_info,
|
||||
)
|
||||
@@ -836,7 +834,7 @@ class ToolNode(RunnableCallable):
|
||||
# Construct ToolRuntime instances at the top level for each tool call
|
||||
tool_runtimes = []
|
||||
for call, cfg in zip(tool_calls, config_list, strict=False):
|
||||
state = self._extract_state(input, cfg)
|
||||
state = self._extract_state(input)
|
||||
tool_runtime = ToolRuntime(
|
||||
state=state,
|
||||
tool_call_id=call["id"],
|
||||
@@ -844,7 +842,6 @@ class ToolNode(RunnableCallable):
|
||||
context=runtime.context,
|
||||
store=runtime.store,
|
||||
stream_writer=runtime.stream_writer,
|
||||
tools=list(self.tools_by_name.values()),
|
||||
execution_info=runtime.execution_info,
|
||||
server_info=runtime.server_info,
|
||||
)
|
||||
@@ -860,30 +857,14 @@ class ToolNode(RunnableCallable):
|
||||
|
||||
def _combine_tool_outputs(
|
||||
self,
|
||||
outputs: list[ToolMessage | Command | list[ToolMessage | Command]],
|
||||
outputs: list[ToolMessage | Command],
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
) -> list[Command | list[ToolMessage] | dict[str, list[ToolMessage]]]:
|
||||
# Flatten list entries from tools that returned multiple items
|
||||
flat_outputs: list[ToolMessage | Command]
|
||||
if any(isinstance(output, list) for output in outputs):
|
||||
flat_outputs = []
|
||||
for output in outputs:
|
||||
if isinstance(output, list):
|
||||
flat_outputs.extend(output)
|
||||
else:
|
||||
flat_outputs.append(output)
|
||||
else:
|
||||
flat_outputs = cast("list[ToolMessage | Command]", outputs)
|
||||
|
||||
# preserve existing behavior for non-command tool outputs for backwards
|
||||
# compatibility
|
||||
if not any(isinstance(output, Command) for output in flat_outputs):
|
||||
if not any(isinstance(output, Command) for output in outputs):
|
||||
# TypedDict, pydantic, dataclass, etc. should all be able to load from dict
|
||||
return (
|
||||
flat_outputs
|
||||
if input_type == "list"
|
||||
else {self._messages_key: flat_outputs}
|
||||
)
|
||||
return outputs if input_type == "list" else {self._messages_key: outputs}
|
||||
|
||||
# LangGraph will automatically handle list of Command and non-command node
|
||||
# updates
|
||||
@@ -893,7 +874,7 @@ class ToolNode(RunnableCallable):
|
||||
|
||||
# combine all parent commands with goto into a single parent command
|
||||
parent_command: Command | None = None
|
||||
for output in flat_outputs:
|
||||
for output in outputs:
|
||||
if isinstance(output, Command):
|
||||
if (
|
||||
output.graph is Command.PARENT
|
||||
@@ -923,7 +904,7 @@ class ToolNode(RunnableCallable):
|
||||
request: ToolCallRequest,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
config: RunnableConfig,
|
||||
) -> ToolMessage | Command | list[Command | ToolMessage]:
|
||||
) -> ToolMessage | Command:
|
||||
"""Execute tool call with configured error handling.
|
||||
|
||||
Args:
|
||||
@@ -932,7 +913,7 @@ class ToolNode(RunnableCallable):
|
||||
config: Runnable configuration.
|
||||
|
||||
Returns:
|
||||
ToolMessage, Command, or list of Command/ToolMessage.
|
||||
ToolMessage or Command.
|
||||
|
||||
Raises:
|
||||
Exception: If tool fails and handle_tool_errors is False.
|
||||
@@ -964,11 +945,6 @@ class ToolNode(RunnableCallable):
|
||||
call["name"], exc, call["args"], filtered_errors
|
||||
) from exc
|
||||
|
||||
# Inside try so validation errors route through _handle_tool_errors
|
||||
return self._normalize_tool_response(
|
||||
response, request.tool_call, input_type
|
||||
)
|
||||
|
||||
# GraphInterrupt is a special exception that will always be raised.
|
||||
# It can be triggered in the following scenarios,
|
||||
# Where GraphInterrupt(GraphBubbleUp) is raised from an `interrupt` invocation
|
||||
@@ -1010,12 +986,23 @@ class ToolNode(RunnableCallable):
|
||||
status="error",
|
||||
)
|
||||
|
||||
# Process successful response
|
||||
if isinstance(response, Command):
|
||||
# Validate Command before returning to handler
|
||||
return self._validate_tool_command(response, request.tool_call, input_type)
|
||||
if isinstance(response, ToolMessage):
|
||||
response.content = cast("str | list", msg_content_output(response.content))
|
||||
return response
|
||||
|
||||
msg = f"Tool {call['name']} returned unexpected type: {type(response)}"
|
||||
raise TypeError(msg)
|
||||
|
||||
def _run_one(
|
||||
self,
|
||||
call: ToolCall,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
tool_runtime: ToolRuntime,
|
||||
) -> ToolMessage | Command | list[Command | ToolMessage]:
|
||||
) -> ToolMessage | Command:
|
||||
"""Execute single tool call with wrap_tool_call wrapper if configured.
|
||||
|
||||
Args:
|
||||
@@ -1070,7 +1057,7 @@ class ToolNode(RunnableCallable):
|
||||
request: ToolCallRequest,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
config: RunnableConfig,
|
||||
) -> ToolMessage | Command | list[Command | ToolMessage]:
|
||||
) -> ToolMessage | Command:
|
||||
"""Execute tool call asynchronously with configured error handling.
|
||||
|
||||
Args:
|
||||
@@ -1079,7 +1066,7 @@ class ToolNode(RunnableCallable):
|
||||
config: Runnable configuration.
|
||||
|
||||
Returns:
|
||||
ToolMessage, Command, or list of Command/ToolMessage.
|
||||
ToolMessage or Command.
|
||||
|
||||
Raises:
|
||||
Exception: If tool fails and handle_tool_errors is False.
|
||||
@@ -1111,11 +1098,6 @@ class ToolNode(RunnableCallable):
|
||||
call["name"], exc, call["args"], filtered_errors
|
||||
) from exc
|
||||
|
||||
# Inside try so validation errors route through _handle_tool_errors
|
||||
return self._normalize_tool_response(
|
||||
response, request.tool_call, input_type
|
||||
)
|
||||
|
||||
# GraphInterrupt is a special exception that will always be raised.
|
||||
# It can be triggered in the following scenarios,
|
||||
# Where GraphInterrupt(GraphBubbleUp) is raised from an `interrupt` invocation
|
||||
@@ -1157,12 +1139,23 @@ class ToolNode(RunnableCallable):
|
||||
status="error",
|
||||
)
|
||||
|
||||
# Process successful response
|
||||
if isinstance(response, Command):
|
||||
# Validate Command before returning to handler
|
||||
return self._validate_tool_command(response, request.tool_call, input_type)
|
||||
if isinstance(response, ToolMessage):
|
||||
response.content = cast("str | list", msg_content_output(response.content))
|
||||
return response
|
||||
|
||||
msg = f"Tool {call['name']} returned unexpected type: {type(response)}"
|
||||
raise TypeError(msg)
|
||||
|
||||
async def _arun_one(
|
||||
self,
|
||||
call: ToolCall,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
tool_runtime: ToolRuntime,
|
||||
) -> ToolMessage | Command | list[Command | ToolMessage]:
|
||||
) -> ToolMessage | Command:
|
||||
"""Execute single tool call asynchronously with awrap_tool_call wrapper if configured.
|
||||
|
||||
Args:
|
||||
@@ -1278,37 +1271,18 @@ class ToolNode(RunnableCallable):
|
||||
return None
|
||||
|
||||
def _extract_state(
|
||||
self,
|
||||
input: list[AnyMessage] | dict[str, Any] | BaseModel,
|
||||
config: RunnableConfig,
|
||||
self, input: list[AnyMessage] | dict[str, Any] | BaseModel
|
||||
) -> list[AnyMessage] | dict[str, Any] | BaseModel:
|
||||
"""Extract state from input.
|
||||
"""Extract state from input, handling ToolCallWithContext if present.
|
||||
|
||||
Three input shapes:
|
||||
Args:
|
||||
input: The input which may be raw state or ToolCallWithContext.
|
||||
|
||||
- `ToolCallWithContext` dict — legacy Send payload carrying an inlined
|
||||
state snapshot; return `input["state"]`.
|
||||
- list of `ToolCall` dicts — new Send payload with no inlined state;
|
||||
hydrate state from channels via `CONFIG_KEY_READ`.
|
||||
- regular graph state (dict/list/BaseModel) — return `input` as-is.
|
||||
Returns:
|
||||
The actual state to pass to wrap_tool_call wrappers.
|
||||
"""
|
||||
if isinstance(input, dict) and input.get("__type") == "tool_call_with_context":
|
||||
return input["state"]
|
||||
if (
|
||||
isinstance(input, list)
|
||||
and input
|
||||
and isinstance(input[-1], dict)
|
||||
and input[-1].get("type") == "tool_call"
|
||||
):
|
||||
read = config.get(CONF, {}).get(CONFIG_KEY_READ)
|
||||
if read is None:
|
||||
return {}
|
||||
# Pregel installs CONFIG_KEY_READ as
|
||||
# `functools.partial(local_read, scratchpad, channels, managed, task)`.
|
||||
# Match the previous inlined-state contract by reading channels only;
|
||||
# managed values have their own injection path (`ToolRuntime.context`).
|
||||
channels = read.args[1]
|
||||
return cast("dict[str, Any]", read(list(channels), True))
|
||||
return input
|
||||
|
||||
def _inject_tool_args(
|
||||
@@ -1428,84 +1402,11 @@ class ToolNode(RunnableCallable):
|
||||
tool_call_copy["args"] = {**stripped_args, **injected_args}
|
||||
return tool_call_copy
|
||||
|
||||
def _normalize_tool_response(
|
||||
self,
|
||||
response: Any,
|
||||
tool_call: ToolCall,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
) -> ToolMessage | Command | list[Command | ToolMessage]:
|
||||
"""Validate and normalize a tool's raw return value."""
|
||||
if isinstance(response, Command):
|
||||
return self._validate_tool_command(response, tool_call, input_type)
|
||||
if isinstance(response, ToolMessage):
|
||||
response.content = cast("str | list", msg_content_output(response.content))
|
||||
return response
|
||||
if isinstance(response, list):
|
||||
if all(isinstance(r, (Command, ToolMessage)) for r in response):
|
||||
return self._validate_tool_command_list(response, tool_call, input_type)
|
||||
msg = (
|
||||
f"Tool {tool_call['name']} returned a list with invalid element "
|
||||
"types: expected all Command or ToolMessage"
|
||||
)
|
||||
raise TypeError(msg)
|
||||
msg = f"Tool {tool_call['name']} returned unexpected type: {type(response)}"
|
||||
raise TypeError(msg)
|
||||
|
||||
def _validate_tool_command_list(
|
||||
self,
|
||||
response: list[Command | ToolMessage],
|
||||
tool_call: ToolCall,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
) -> list[Command | ToolMessage]:
|
||||
"""Validate a list of Command/ToolMessage returned by a single tool call.
|
||||
|
||||
Requires exactly one terminating ToolMessage (matching the outer tool_call_id)
|
||||
across the list — either as a top-level element or nested in a
|
||||
Command.update["messages"].
|
||||
"""
|
||||
expected_id = tool_call["id"]
|
||||
|
||||
terminator_count = 0
|
||||
for item in response:
|
||||
if isinstance(item, ToolMessage):
|
||||
if item.tool_call_id == expected_id:
|
||||
terminator_count += 1
|
||||
elif isinstance(item, Command) and isinstance(item.update, dict):
|
||||
for msg in item.update.get(self._messages_key, []):
|
||||
if isinstance(msg, ToolMessage) and msg.tool_call_id == expected_id:
|
||||
terminator_count += 1
|
||||
|
||||
if terminator_count != 1:
|
||||
msg = (
|
||||
f"Tool {tool_call['name']} returned a list with "
|
||||
f"{terminator_count} messages bound to tool_call_id "
|
||||
f"{expected_id!r}; expected exactly one terminating ToolMessage."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
|
||||
# Per-Command normalization still runs, but the list-level count above
|
||||
# already guarantees exactly one terminator, so individual Commands may
|
||||
# lack one.
|
||||
validated: list[Command | ToolMessage] = []
|
||||
for item in response:
|
||||
if isinstance(item, Command):
|
||||
validated.append(
|
||||
self._validate_tool_command(
|
||||
item, tool_call, input_type, require_terminator=False
|
||||
)
|
||||
)
|
||||
else:
|
||||
item.content = cast("str | list", msg_content_output(item.content))
|
||||
validated.append(item)
|
||||
return validated
|
||||
|
||||
def _validate_tool_command(
|
||||
self,
|
||||
command: Command,
|
||||
call: ToolCall,
|
||||
input_type: Literal["list", "dict", "tool_calls"],
|
||||
*,
|
||||
require_terminator: bool = True,
|
||||
) -> Command:
|
||||
if isinstance(command.update, dict):
|
||||
# input type is dict when ToolNode is invoked with a dict input
|
||||
@@ -1555,11 +1456,7 @@ class ToolNode(RunnableCallable):
|
||||
|
||||
# validate that we always have a ToolMessage matching the tool call in
|
||||
# Command.update if command is sent to the CURRENT graph
|
||||
if (
|
||||
require_terminator
|
||||
and updated_command.graph is None
|
||||
and not has_matching_tool_message
|
||||
):
|
||||
if updated_command.graph is None and not has_matching_tool_message:
|
||||
example_update = (
|
||||
'`Command(update={"messages": '
|
||||
'[ToolMessage("Success", tool_call_id=tool_call_id), ...]}, ...)`'
|
||||
@@ -1679,7 +1576,6 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
|
||||
- `context`: Runtime context (shared with `Runtime`)
|
||||
- `store`: `BaseStore` instance for persistent storage (shared with `Runtime`)
|
||||
- `stream_writer`: `StreamWriter` for streaming output (shared with `Runtime`)
|
||||
- `tools`: List of all available `BaseTool` instances
|
||||
|
||||
No `Annotated` wrapper is needed - just use `runtime: ToolRuntime`
|
||||
as a parameter.
|
||||
@@ -1722,7 +1618,6 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
|
||||
context: ContextT
|
||||
config: RunnableConfig
|
||||
stream_writer: StreamWriter
|
||||
tools: list[BaseTool]
|
||||
tool_call_id: str | None
|
||||
store: BaseStore | None
|
||||
execution_info: ExecutionInfo | None = None
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.12"
|
||||
version = "1.0.9"
|
||||
description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
@@ -25,7 +25,7 @@ classifiers = [
|
||||
]
|
||||
dependencies = [
|
||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||
"langchain-core>=1.3.1",
|
||||
"langchain-core>=1.0.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
|
||||
@@ -4,7 +4,7 @@ This tests the fix for https://github.com/langchain-ai/langchain/issues/35585
|
||||
|
||||
When using InjectedState(<field>) on a tool parameter, and the referenced field is
|
||||
declared as NotRequired in the custom state schema, the ToolNode should gracefully
|
||||
handle missing fields by injecting None instead of raising KeyError.
|
||||
handle missing fields without raising KeyError so the tool's default can apply.
|
||||
"""
|
||||
|
||||
import sys
|
||||
@@ -45,6 +45,14 @@ def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> st
|
||||
return f"It's always sunny in {city}!"
|
||||
|
||||
|
||||
@tool
|
||||
def get_weather_with_default(
|
||||
city: Annotated[str, InjectedState("city")] = "Boston",
|
||||
) -> str:
|
||||
"""Get weather for a given city, defaulting when state omits the field."""
|
||||
return f"It's always sunny in {city}!"
|
||||
|
||||
|
||||
def _create_mock_runtime(
|
||||
state: dict | None = None,
|
||||
store=None,
|
||||
@@ -69,7 +77,6 @@ def _create_config_with_runtime(store=None, state=None):
|
||||
context={},
|
||||
store=store,
|
||||
stream_writer=None,
|
||||
tools=[],
|
||||
tool_call_id="test_id",
|
||||
)
|
||||
return {
|
||||
@@ -85,7 +92,7 @@ def _create_config_with_runtime(store=None, state=None):
|
||||
reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
|
||||
)
|
||||
def test_injected_state_not_required_field_missing_injects_none():
|
||||
"""Test that InjectedState with NotRequired field injects None when field is missing.
|
||||
"""Test that missing optional InjectedState leaves the tool default in place.
|
||||
|
||||
This verifies the fix for https://github.com/langchain-ai/langchain/issues/35585
|
||||
"""
|
||||
@@ -115,6 +122,37 @@ def test_injected_state_not_required_field_missing_injects_none():
|
||||
assert "No city provided" in tool_msg.content
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
|
||||
)
|
||||
def test_injected_state_not_required_field_missing_preserves_tool_default():
|
||||
"""Test that missing optional InjectedState preserves a non-None tool default."""
|
||||
tool_node = ToolNode([get_weather_with_default])
|
||||
|
||||
tool_call = {
|
||||
"name": "get_weather_with_default",
|
||||
"args": {},
|
||||
"id": "call_1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
ai_msg = AIMessage("Let me check the weather", tool_calls=[tool_call])
|
||||
|
||||
state_without_city: CustomAgentStateWithNotRequired = {
|
||||
"messages": [HumanMessage("What's the weather?"), ai_msg],
|
||||
}
|
||||
|
||||
result = tool_node.invoke(
|
||||
state_without_city,
|
||||
config=_create_config_with_runtime(state=state_without_city),
|
||||
)
|
||||
|
||||
assert len(result["messages"]) == 1
|
||||
tool_msg = result["messages"][0]
|
||||
assert isinstance(tool_msg, ToolMessage)
|
||||
assert "Boston" in tool_msg.content
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
|
||||
|
||||
@@ -1320,98 +1320,6 @@ async def test_state_extraction_with_tool_call_with_context_async() -> None:
|
||||
assert "tool_call" not in state_seen[0]
|
||||
|
||||
|
||||
def _config_with_channel_read(
|
||||
channel_values: dict[str, object],
|
||||
store: BaseStore | None = None,
|
||||
) -> RunnableConfig:
|
||||
"""Build a config that mimics `CONFIG_KEY_READ` as Pregel installs it.
|
||||
|
||||
Pregel always installs a `functools.partial(local_read, scratchpad,
|
||||
channels, managed, task)`, and `ToolNode` introspects that partial to
|
||||
learn channel names. The stub matches the shape: partial whose second and
|
||||
third positional args are `channels` and `managed` mappings.
|
||||
"""
|
||||
import functools
|
||||
|
||||
channels_stub = {k: None for k in channel_values}
|
||||
managed_stub: dict[str, object] = {}
|
||||
|
||||
# Shape matches pregel's real partial:
|
||||
# functools.partial(local_read, scratchpad, channels, managed, task)
|
||||
def _read(scratchpad, channels, managed, task, select, fresh): # noqa: ARG001
|
||||
if isinstance(select, str):
|
||||
return channel_values[select]
|
||||
return {k: channel_values[k] for k in select if k in channel_values}
|
||||
|
||||
read = functools.partial(_read, None, channels_stub, managed_stub, None)
|
||||
cfg = _create_config_with_runtime(store)
|
||||
cfg["configurable"]["__pregel_read"] = read
|
||||
return cfg
|
||||
|
||||
|
||||
def test_list_form_send_hydrates_state_from_channel_read() -> None:
|
||||
"""Send('tools', [tool_call]) with no inlined state should hydrate
|
||||
ToolRuntime.state from CONFIG_KEY_READ (full state read)."""
|
||||
state_seen = []
|
||||
|
||||
def state_inspector_handler(
|
||||
request: ToolCallRequest,
|
||||
execute: Callable[[ToolCallRequest], ToolMessage | Command],
|
||||
) -> ToolMessage | Command:
|
||||
state_seen.append(request.state)
|
||||
return execute(request)
|
||||
|
||||
channel_values = {
|
||||
"messages": [AIMessage("from channels")],
|
||||
"files": {"/a.md": "body"},
|
||||
}
|
||||
|
||||
tool_node = ToolNode([add], wrap_tool_call=state_inspector_handler)
|
||||
|
||||
tool_call: ToolCall = {
|
||||
"name": "add",
|
||||
"args": {"a": 1, "b": 2},
|
||||
"id": "call_1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
|
||||
tool_node.invoke([tool_call], config=_config_with_channel_read(channel_values))
|
||||
|
||||
assert len(state_seen) == 1
|
||||
got = state_seen[0]
|
||||
assert got == channel_values
|
||||
assert "messages" in got and "files" in got
|
||||
|
||||
|
||||
async def test_list_form_send_hydrates_state_async() -> None:
|
||||
state_seen = []
|
||||
|
||||
def state_inspector_handler(
|
||||
request: ToolCallRequest,
|
||||
execute: Callable[[ToolCallRequest], ToolMessage | Command],
|
||||
) -> ToolMessage | Command:
|
||||
state_seen.append(request.state)
|
||||
return execute(request)
|
||||
|
||||
channel_values = {"messages": [AIMessage("from channels")], "files": {}}
|
||||
|
||||
tool_node = ToolNode([add], wrap_tool_call=state_inspector_handler)
|
||||
|
||||
tool_call: ToolCall = {
|
||||
"name": "add",
|
||||
"args": {"a": 1, "b": 2},
|
||||
"id": "call_1",
|
||||
"type": "tool_call",
|
||||
}
|
||||
|
||||
await tool_node.ainvoke(
|
||||
[tool_call], config=_config_with_channel_read(channel_values)
|
||||
)
|
||||
|
||||
assert len(state_seen) == 1
|
||||
assert state_seen[0] == channel_values
|
||||
|
||||
|
||||
def test_tool_call_request_is_frozen() -> None:
|
||||
"""Test that ToolCallRequest raises deprecation warnings on direct attribute reassignment."""
|
||||
tool_call: ToolCall = {"name": "add", "args": {"a": 1, "b": 2}, "id": "call_1"}
|
||||
|
||||
@@ -2016,8 +2016,8 @@ async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async()
|
||||
assert tool_message.tool_call_id == "call_dynamic_2"
|
||||
|
||||
|
||||
def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
"""Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""
|
||||
def test_tool_runtime_forwards_execution_info_and_server_info() -> None:
|
||||
"""Test that execution_info and server_info are forwarded from Runtime to ToolRuntime."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
|
||||
exec_info = ExecutionInfo(
|
||||
@@ -2043,15 +2043,9 @@ def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
"""Tool that captures runtime info."""
|
||||
captured["execution_info"] = runtime.execution_info
|
||||
captured["server_info"] = runtime.server_info
|
||||
captured["tools"] = runtime.tools
|
||||
return "ok"
|
||||
|
||||
@dec_tool
|
||||
def other_tool(y: int) -> str:
|
||||
"""Another tool available to the runtime."""
|
||||
return str(y)
|
||||
|
||||
node = ToolNode([info_tool, other_tool])
|
||||
node = ToolNode([info_tool])
|
||||
tool_call = {
|
||||
"name": "info_tool",
|
||||
"args": {"x": 1},
|
||||
@@ -2060,21 +2054,17 @@ def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
||||
result = node.invoke({"messages": [msg]}, config=config)
|
||||
node.invoke({"messages": [msg]}, config=config)
|
||||
|
||||
assert result["messages"][-1].content == "ok"
|
||||
assert captured["execution_info"] is exec_info
|
||||
assert captured["execution_info"].thread_id == "t-1"
|
||||
assert captured["execution_info"].task_id == "tk-1"
|
||||
assert captured["server_info"] is server_info
|
||||
assert captured["server_info"].assistant_id == "asst-1"
|
||||
assert [tool.name for tool in captured["tools"]] == ["info_tool", "other_tool"]
|
||||
|
||||
|
||||
async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async() -> (
|
||||
None
|
||||
):
|
||||
"""Test that execution_info, server_info, and tools are forwarded in async path."""
|
||||
async def test_tool_runtime_forwards_execution_info_and_server_info_async() -> None:
|
||||
"""Test that execution_info and server_info are forwarded in async path."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
|
||||
exec_info = ExecutionInfo(
|
||||
@@ -2100,15 +2090,9 @@ async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async(
|
||||
"""Async tool that captures runtime info."""
|
||||
captured["execution_info"] = runtime.execution_info
|
||||
captured["server_info"] = runtime.server_info
|
||||
captured["tools"] = runtime.tools
|
||||
return "ok"
|
||||
|
||||
@dec_tool
|
||||
async def other_tool_async(y: int) -> str:
|
||||
"""Another async tool available to the runtime."""
|
||||
return str(y)
|
||||
|
||||
node = ToolNode([info_tool_async, other_tool_async])
|
||||
node = ToolNode([info_tool_async])
|
||||
tool_call = {
|
||||
"name": "info_tool_async",
|
||||
"args": {"x": 1},
|
||||
@@ -2117,17 +2101,12 @@ async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async(
|
||||
}
|
||||
msg = AIMessage("", tool_calls=[tool_call])
|
||||
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
|
||||
result = await node.ainvoke({"messages": [msg]}, config=config)
|
||||
await node.ainvoke({"messages": [msg]}, config=config)
|
||||
|
||||
assert result["messages"][-1].content == "ok"
|
||||
assert captured["execution_info"] is exec_info
|
||||
assert captured["execution_info"].thread_id == "t-2"
|
||||
assert captured["server_info"] is server_info
|
||||
assert captured["server_info"].graph_id == "graph-2"
|
||||
assert [tool.name for tool in captured["tools"]] == [
|
||||
"info_tool_async",
|
||||
"other_tool_async",
|
||||
]
|
||||
|
||||
|
||||
# --- InjectedToolArg security tests ---
|
||||
@@ -2223,195 +2202,3 @@ def test_tool_node_injected_state_overwrites_llm_value() -> None:
|
||||
)
|
||||
tool_message = result["messages"][-1]
|
||||
assert tool_message.content == "PUBLIC_DATA"
|
||||
|
||||
|
||||
class _ReturningTool(BaseTool):
|
||||
"""A tool that returns a configured value verbatim."""
|
||||
|
||||
name: str = "list_tool"
|
||||
description: str = "Returns a configured value"
|
||||
return_value: Any = None
|
||||
|
||||
def _run(self, **kwargs: Any) -> Any:
|
||||
return self.return_value
|
||||
|
||||
async def _arun(self, **kwargs: Any) -> Any:
|
||||
return self.return_value
|
||||
|
||||
|
||||
def _list_tool_call(outer_id: str = "call-1") -> dict[str, Any]:
|
||||
return {"name": "list_tool", "args": {}, "id": outer_id, "type": "tool_call"}
|
||||
|
||||
|
||||
def _invoke_returning(
|
||||
return_value: Any,
|
||||
*,
|
||||
outer_id: str = "call-1",
|
||||
handle_tool_errors: bool = True,
|
||||
) -> Any:
|
||||
node = ToolNode(
|
||||
[_ReturningTool(return_value=return_value)],
|
||||
handle_tool_errors=handle_tool_errors,
|
||||
)
|
||||
return node.invoke(
|
||||
{"messages": [AIMessage("", tool_calls=[_list_tool_call(outer_id)])]},
|
||||
config=_create_config_with_runtime(),
|
||||
)
|
||||
|
||||
|
||||
def test_tool_node_list_return_command_and_tool_message() -> None:
|
||||
"""Valid: tool returns [Command(update={...}), ToolMessage(...)]."""
|
||||
outer_id = "call-1"
|
||||
result = _invoke_returning(
|
||||
[
|
||||
Command(update={"foo": "bar"}),
|
||||
ToolMessage(content="done", tool_call_id=outer_id),
|
||||
]
|
||||
)
|
||||
assert isinstance(result, list)
|
||||
commands = [r for r in result if isinstance(r, Command)]
|
||||
assert len(commands) == 1
|
||||
assert commands[0].update == {"foo": "bar"}
|
||||
non_commands = [r for r in result if not isinstance(r, Command)]
|
||||
assert len(non_commands) == 1
|
||||
assert isinstance(non_commands[0], dict)
|
||||
msgs = non_commands[0]["messages"]
|
||||
assert len(msgs) == 1
|
||||
assert isinstance(msgs[0], ToolMessage)
|
||||
assert msgs[0].content == "done"
|
||||
assert msgs[0].tool_call_id == outer_id
|
||||
|
||||
|
||||
def test_tool_node_list_return_nested_terminator() -> None:
|
||||
"""Valid: terminator nested inside Command.update['messages']."""
|
||||
outer_id = "call-1"
|
||||
result = _invoke_returning(
|
||||
[
|
||||
Command(update={"foo": "bar"}),
|
||||
Command(
|
||||
update={
|
||||
"messages": [ToolMessage(content="done", tool_call_id=outer_id)]
|
||||
}
|
||||
),
|
||||
]
|
||||
)
|
||||
assert isinstance(result, list)
|
||||
commands = [r for r in result if isinstance(r, Command)]
|
||||
assert len(commands) == 2
|
||||
updates = [c.update for c in commands]
|
||||
assert {"foo": "bar"} in updates
|
||||
msgs_update = next(u for u in updates if "messages" in (u or {}))
|
||||
assert any(
|
||||
isinstance(m, ToolMessage) and m.tool_call_id == outer_id
|
||||
for m in msgs_update["messages"]
|
||||
)
|
||||
|
||||
|
||||
def test_tool_node_list_return_parent_goto_with_terminator() -> None:
|
||||
"""Valid: [Command(graph=PARENT, goto=[Send(...)]), ToolMessage(...)]."""
|
||||
outer_id = "call-1"
|
||||
result = _invoke_returning(
|
||||
[
|
||||
Command(graph=Command.PARENT, goto=[Send("child", {})]),
|
||||
ToolMessage(content="ok", tool_call_id=outer_id),
|
||||
]
|
||||
)
|
||||
assert isinstance(result, list)
|
||||
parent_cmds = [
|
||||
r for r in result if isinstance(r, Command) and r.graph is Command.PARENT
|
||||
]
|
||||
assert len(parent_cmds) == 1
|
||||
assert isinstance(parent_cmds[0].goto, list)
|
||||
assert any(isinstance(s, Send) for s in parent_cmds[0].goto)
|
||||
non_commands = [r for r in result if not isinstance(r, Command)]
|
||||
assert len(non_commands) == 1
|
||||
|
||||
|
||||
def test_tool_node_list_return_no_terminator_raises() -> None:
|
||||
"""Invalid: list with no terminating ToolMessage."""
|
||||
with pytest.raises(ValueError, match="0 messages bound to tool_call_id"):
|
||||
_invoke_returning([Command(update={"foo": "bar"})], handle_tool_errors=False)
|
||||
|
||||
|
||||
def test_tool_node_list_return_multiple_terminators_raises() -> None:
|
||||
"""Invalid: list with two terminating ToolMessages."""
|
||||
outer_id = "call-1"
|
||||
with pytest.raises(ValueError, match="2 messages bound to tool_call_id"):
|
||||
_invoke_returning(
|
||||
[
|
||||
ToolMessage(content="a", tool_call_id=outer_id),
|
||||
ToolMessage(content="b", tool_call_id=outer_id),
|
||||
],
|
||||
handle_tool_errors=False,
|
||||
)
|
||||
|
||||
|
||||
def test_tool_node_list_return_validation_error_handled() -> None:
|
||||
"""handle_tool_errors=True converts validation errors to an error ToolMessage."""
|
||||
result = _invoke_returning([Command(update={"foo": "bar"})])
|
||||
assert isinstance(result, dict)
|
||||
msg = result["messages"][0]
|
||||
assert isinstance(msg, ToolMessage)
|
||||
assert msg.status == "error"
|
||||
assert "0 messages bound to tool_call_id" in msg.content
|
||||
|
||||
|
||||
async def test_tool_node_list_return_async_smoke() -> None:
|
||||
"""Async path parallels sync for the happy case."""
|
||||
outer_id = "call-1"
|
||||
node = ToolNode(
|
||||
[
|
||||
_ReturningTool(
|
||||
return_value=[
|
||||
Command(update={"foo": "bar"}),
|
||||
ToolMessage(content="done", tool_call_id=outer_id),
|
||||
]
|
||||
)
|
||||
]
|
||||
)
|
||||
result = await node.ainvoke(
|
||||
{"messages": [AIMessage("", tool_calls=[_list_tool_call(outer_id)])]},
|
||||
config=_create_config_with_runtime(),
|
||||
)
|
||||
assert isinstance(result, list)
|
||||
commands = [r for r in result if isinstance(r, Command)]
|
||||
assert len(commands) == 1 and commands[0].update == {"foo": "bar"}
|
||||
|
||||
|
||||
def test_tool_node_list_return_mixed_with_regular_tool() -> None:
|
||||
"""List-returning tool and a regular tool dispatched from the same AIMessage."""
|
||||
list_tool_id = "call-list"
|
||||
regular_tool_id = "call-regular"
|
||||
list_tool = _ReturningTool(
|
||||
return_value=[
|
||||
Command(update={"foo": "bar"}),
|
||||
ToolMessage(content="list done", tool_call_id=list_tool_id),
|
||||
]
|
||||
)
|
||||
|
||||
def regular_tool(x: int) -> str:
|
||||
"""A normal tool."""
|
||||
return f"regular: {x}"
|
||||
|
||||
tool_calls = [
|
||||
{"name": "list_tool", "args": {}, "id": list_tool_id, "type": "tool_call"},
|
||||
{
|
||||
"name": "regular_tool",
|
||||
"args": {"x": 7},
|
||||
"id": regular_tool_id,
|
||||
"type": "tool_call",
|
||||
},
|
||||
]
|
||||
node = ToolNode([list_tool, regular_tool])
|
||||
result = node.invoke(
|
||||
{"messages": [AIMessage("", tool_calls=tool_calls)]},
|
||||
config=_create_config_with_runtime(),
|
||||
)
|
||||
assert isinstance(result, list)
|
||||
commands = [r for r in result if isinstance(r, Command)]
|
||||
assert len(commands) == 1
|
||||
assert commands[0].update == {"foo": "bar"}
|
||||
all_msgs = [m for r in result if isinstance(r, dict) for m in r["messages"]]
|
||||
tool_call_ids = {m.tool_call_id for m in all_msgs}
|
||||
assert list_tool_id in tool_call_ids
|
||||
assert regular_tool_id in tool_call_ids
|
||||
|
||||
Generated
+8
-8
@@ -249,7 +249,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.0a2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -261,14 +261,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.1.7a2"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -281,7 +281,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "." },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
@@ -352,7 +352,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -490,7 +490,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.12"
|
||||
version = "1.0.9"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -535,7 +535,7 @@ test = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.1" },
|
||||
{ name = "langchain-core", specifier = ">=1.0.0" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
]
|
||||
|
||||
|
||||
Generated
+8
-8
@@ -262,7 +262,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.0a2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -274,14 +274,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.1.7a2"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -294,7 +294,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "." },
|
||||
@@ -365,7 +365,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.0.2"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -413,7 +413,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.12"
|
||||
version = "1.0.9"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -422,7 +422,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.1" },
|
||||
{ name = "langchain-core", specifier = ">=1.0.0" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user