mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-26 17:42:24 +02:00
Compare commits
90
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/
|
||||
|
||||
@@ -4,15 +4,17 @@ import threading
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -23,7 +25,11 @@ from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_COMBINED_SQL,
|
||||
BasePostgresSaver,
|
||||
_DeltaCombinedRow,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
|
||||
Conn = _internal.Conn # For backward compatibility
|
||||
@@ -430,6 +436,48 @@ class PostgresSaver(BasePostgresSaver):
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def _get_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
|
||||
|
||||
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
|
||||
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
|
||||
single roundtrip; the ancestor walk runs in Python.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
# Caller didn't specify a target — resolve to the latest
|
||||
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
|
||||
# returns the newest; its config carries the resolved id.
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
with self._cursor() as cur:
|
||||
cur.execute(
|
||||
SELECT_DELTA_COMBINED_SQL,
|
||||
(
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
return self._build_delta_channel_writes_history(
|
||||
channel=channel,
|
||||
target_id=checkpoint_id,
|
||||
rows=cast("list[_DeltaCombinedRow]", rows),
|
||||
)
|
||||
|
||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
|
||||
@@ -4,15 +4,17 @@ import asyncio
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -23,7 +25,11 @@ from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_COMBINED_SQL,
|
||||
BasePostgresSaver,
|
||||
_DeltaCombinedRow,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
|
||||
Conn = _ainternal.Conn # For backward compatibility
|
||||
@@ -391,6 +397,46 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _aget_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
|
||||
|
||||
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
|
||||
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
|
||||
single roundtrip; rows are assembled by the shared pure helper on
|
||||
`BasePostgresSaver`.
|
||||
"""
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
target = await self.aget_tuple(config)
|
||||
if target is None:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
async with self._cursor() as cur:
|
||||
await cur.execute(
|
||||
SELECT_DELTA_COMBINED_SQL,
|
||||
(
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
),
|
||||
)
|
||||
rows = await cur.fetchall()
|
||||
return self._build_delta_channel_writes_history(
|
||||
channel=channel,
|
||||
target_id=checkpoint_id,
|
||||
rows=cast("list[_DeltaCombinedRow]", rows),
|
||||
)
|
||||
|
||||
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
|
||||
@@ -4,13 +4,16 @@ import random
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from importlib.metadata import version as get_version
|
||||
from typing import Any, cast
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
PendingWrite,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
@@ -153,6 +156,62 @@ INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
"""
|
||||
|
||||
|
||||
class _DeltaCombinedRow(TypedDict, total=False):
|
||||
"""One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables).
|
||||
|
||||
Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are
|
||||
relevant for that kind; irrelevant columns are NULL and typed as `None`.
|
||||
"""
|
||||
|
||||
_kind: str # always present: "p", "w", or "b"
|
||||
# checkpoint row ("p")
|
||||
checkpoint_id: str | None
|
||||
parent_checkpoint_id: str | None
|
||||
ver: str | None
|
||||
# write / blob rows ("w", "b")
|
||||
type: str | None
|
||||
blob: bytes | None
|
||||
# write row only ("w")
|
||||
task_id: str | None
|
||||
idx: int | None
|
||||
# blob row only ("b")
|
||||
version: str | None
|
||||
|
||||
|
||||
# DeltaChannel reconstruction: one UNION ALL query fetches checkpoints,
|
||||
# writes, and blobs for `channel` in one roundtrip; the ancestor walk runs
|
||||
# in Python in `_build_delta_channel_writes_history`.
|
||||
#
|
||||
# Parameter order: (channel, thread_id, checkpoint_ns,
|
||||
# thread_id, checkpoint_ns, channel,
|
||||
# thread_id, checkpoint_ns, channel)
|
||||
SELECT_DELTA_COMBINED_SQL = """
|
||||
SELECT 'p'::text AS _kind,
|
||||
checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver,
|
||||
NULL::text AS type,
|
||||
NULL::bytea AS blob,
|
||||
NULL::text AS task_id,
|
||||
NULL::int AS idx,
|
||||
NULL::text AS version
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
UNION ALL
|
||||
SELECT 'w',
|
||||
checkpoint_id, NULL, NULL,
|
||||
type, blob, task_id, idx, NULL
|
||||
FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
|
||||
UNION ALL
|
||||
SELECT 'b',
|
||||
NULL, NULL, NULL,
|
||||
type, blob, NULL, NULL, version
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
|
||||
"""
|
||||
|
||||
|
||||
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
SELECT_SQL = SELECT_SQL
|
||||
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
|
||||
@@ -195,6 +254,83 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
if t.decode() != "empty"
|
||||
}
|
||||
|
||||
def _build_delta_channel_writes_history(
|
||||
self,
|
||||
*,
|
||||
channel: str,
|
||||
target_id: str,
|
||||
rows: Sequence[_DeltaCombinedRow],
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Reconstruct one delta channel's history from the combined UNION ALL rows.
|
||||
|
||||
Pure data transform shared by sync (`PostgresSaver`) and async
|
||||
(`AsyncPostgresSaver`); both paths run `SELECT_DELTA_COMBINED_SQL`
|
||||
and feed the tagged rows here.
|
||||
|
||||
Walk is newest → oldest from the target's parent. A non-sentinel
|
||||
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
|
||||
walk and is returned as the seed so replay starts from it.
|
||||
|
||||
Writes stored at `target_id` itself are pending writes for the next
|
||||
step and are excluded — the walk begins at the target's parent.
|
||||
"""
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_of: dict[str, str | None] = {}
|
||||
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
|
||||
blob_by_ver: dict[str, tuple[str, bytes]] = {}
|
||||
|
||||
for r in rows:
|
||||
kind = r["_kind"]
|
||||
if kind == "p":
|
||||
cid = cast(str, r["checkpoint_id"])
|
||||
parent_of[cid] = r["parent_checkpoint_id"]
|
||||
ver_of[cid] = r["ver"]
|
||||
elif kind == "w":
|
||||
cid = cast(str, r["checkpoint_id"])
|
||||
writes_by_cid.setdefault(cid, []).append(
|
||||
cast(
|
||||
"tuple[str, bytes, str, int]",
|
||||
(r["type"], r["blob"], r["task_id"], r["idx"]),
|
||||
)
|
||||
)
|
||||
else: # kind == "b"
|
||||
blob_by_ver[cast(str, r["version"])] = cast(
|
||||
"tuple[str, bytes]", (r["type"], r["blob"])
|
||||
)
|
||||
|
||||
# newest write first per ancestor (task_id DESC, idx DESC)
|
||||
for ws in writes_by_cid.values():
|
||||
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
|
||||
|
||||
ancestors: list[str] = []
|
||||
cur_cid: str | None = parent_of.get(target_id)
|
||||
while cur_cid is not None:
|
||||
ancestors.append(cur_cid)
|
||||
cur_cid = parent_of.get(cur_cid)
|
||||
if not ancestors:
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
|
||||
|
||||
collected: list[PendingWrite] = [] # newest first; reversed at the end
|
||||
for cid in ancestors:
|
||||
# Collect writes first — they encode the transition FROM this
|
||||
# ancestor's state to its child's and must be included even if
|
||||
# this ancestor is also the seed checkpoint.
|
||||
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
|
||||
val = self.serde.loads_typed((type_tag, write_blob))
|
||||
collected.append((task_id, channel, val))
|
||||
# Then check seed terminator.
|
||||
ver = ver_of.get(cid)
|
||||
if ver is not None:
|
||||
seed_blob = blob_by_ver.get(ver)
|
||||
if seed_blob is not None and seed_blob[0] != "empty":
|
||||
blob_value = self.serde.loads_typed(seed_blob)
|
||||
if blob_value is not DELTA_SENTINEL:
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=blob_value, writes=collected)
|
||||
|
||||
collected.reverse() # oldest → newest
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
thread_id: str,
|
||||
|
||||
@@ -12,7 +12,7 @@ readme = "README.md"
|
||||
license = "MIT"
|
||||
license-files = ['LICENSE']
|
||||
dependencies = [
|
||||
"langgraph-checkpoint>=2.1.2,<5.0.0",
|
||||
"langgraph-checkpoint>=4.0.3,<5.0.0",
|
||||
"orjson>=3.11.5",
|
||||
"psycopg>=3.2.0",
|
||||
"psycopg-pool>=3.2.0",
|
||||
|
||||
@@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values(
|
||||
|
||||
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
||||
|
||||
def patched_load_checkpoint_tuple(value):
|
||||
async def patched_load_checkpoint_tuple(value):
|
||||
value["checkpoint"].pop("channel_values", None)
|
||||
return load_checkpoint_tuple(value)
|
||||
return await load_checkpoint_tuple(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
||||
@@ -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 _messages_delta_reducer
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
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" },
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import copy
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
|
||||
from typing import ( # noqa: UP035
|
||||
from typing import (
|
||||
Any,
|
||||
Generic,
|
||||
Literal,
|
||||
@@ -18,6 +18,9 @@ from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
DELTA_SENTINEL as DELTA_SENTINEL,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
ERROR,
|
||||
INTERRUPT,
|
||||
@@ -28,6 +31,8 @@ from langgraph.checkpoint.serde.types import (
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
PendingWrite = tuple[str, str, Any]
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -119,6 +124,30 @@ class CheckpointTuple(NamedTuple):
|
||||
pending_writes: list[PendingWrite] | None = None
|
||||
|
||||
|
||||
class _ChannelWritesHistory(NamedTuple):
|
||||
"""Result of `BaseCheckpointSaver._get_channel_writes_history`.
|
||||
|
||||
Storage-level view of what one channel wrote across the ancestor chain
|
||||
of a target checkpoint:
|
||||
|
||||
* `seed` — the nearest ancestor's stored blob value for this channel,
|
||||
or `DELTA_SENTINEL` if the walk reached the root without finding a
|
||||
stored value. A non-sentinel seed typically indicates a pre-delta
|
||||
snapshot preserved across a channel-type migration (e.g.
|
||||
`BinaryOperatorAggregate` storage extended under `DeltaChannel`).
|
||||
* `writes` — on-path deltas oldest→newest, one `PendingWrite` per
|
||||
step that wrote to this channel. Writes stored at the target
|
||||
checkpoint itself are pending for the next super-step and are
|
||||
excluded.
|
||||
|
||||
Experimental: method surface may change; the NamedTuple shape is the
|
||||
contract.
|
||||
"""
|
||||
|
||||
seed: Any
|
||||
writes: list[PendingWrite]
|
||||
|
||||
|
||||
class BaseCheckpointSaver(Generic[V]):
|
||||
"""Base class for creating a graph checkpointer.
|
||||
|
||||
@@ -457,6 +486,104 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Pure storage read used by `_get_channel_writes_history`.
|
||||
|
||||
Must return the same value as `get_tuple` but must NOT trigger channel
|
||||
reconstruction; otherwise the channel-hydration path would re-enter
|
||||
`_get_channel_writes_history`. Override only if `get_tuple` itself
|
||||
performs channel hydration.
|
||||
"""
|
||||
return self.get_tuple(config)
|
||||
|
||||
async def _aget_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Async version of `_get_tuple_raw`. See docstring there."""
|
||||
return await self.aget_tuple(config)
|
||||
|
||||
def _get_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""**Experimental.** Query one channel's writes along the parent chain.
|
||||
|
||||
Storage-level query, not channel semantics: returns `(seed, writes)`
|
||||
reflecting what storage knows about a single channel across the
|
||||
ancestor chain of the target checkpoint identified by `config`.
|
||||
|
||||
* `writes` — on-path deltas oldest→newest as `PendingWrite` tuples.
|
||||
Writes stored at the target `checkpoint_id` itself are pending
|
||||
for the next super-step and are excluded.
|
||||
* `seed` — the nearest ancestor's stored blob value for this
|
||||
channel; `DELTA_SENTINEL` if the walk reached the root without
|
||||
finding a stored value. A non-sentinel seed typically indicates
|
||||
a pre-delta snapshot preserved across a channel-type migration.
|
||||
|
||||
Walks the **parent chain** (not `list(before=...)`): for forked
|
||||
threads, only on-path ancestors contribute.
|
||||
|
||||
Reference implementation walks `get_tuple` + `parent_config`,
|
||||
inspecting each ancestor's `channel_values[channel]` for the seed
|
||||
terminator. Savers with direct storage access (`InMemorySaver`,
|
||||
`PostgresSaver`) override for performance; the return contract is
|
||||
fixed here.
|
||||
|
||||
Underscore-prefixed because the method surface is experimental.
|
||||
"""
|
||||
collected: list[PendingWrite] = [] # newest first; reversed at the end
|
||||
target_tuple = self._get_tuple_raw(config)
|
||||
cursor_config: RunnableConfig | None = (
|
||||
target_tuple.parent_config if target_tuple else None
|
||||
)
|
||||
while cursor_config is not None:
|
||||
tup = self._get_tuple_raw(cursor_config)
|
||||
if tup is None:
|
||||
break
|
||||
# Collect this ancestor's writes FIRST — they encode the
|
||||
# transition from this ancestor's state to its child's, so
|
||||
# they must be included whether or not this ancestor is the
|
||||
# seed terminator.
|
||||
if tup.pending_writes:
|
||||
# Within a superstep, pending_writes are oldest→newest;
|
||||
# reverse to scan newest-first.
|
||||
for write in reversed(tup.pending_writes):
|
||||
if write[1] != channel:
|
||||
continue
|
||||
collected.append(write)
|
||||
# Seed terminator: any non-sentinel blob on an ancestor
|
||||
# establishes the reconstruction base. Stop here.
|
||||
ancestor_value = tup.checkpoint["channel_values"].get(channel)
|
||||
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=ancestor_value, writes=collected)
|
||||
cursor_config = tup.parent_config
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
|
||||
|
||||
async def _aget_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
"""Async version of `_get_channel_writes_history`. See docstring there."""
|
||||
collected: list[PendingWrite] = []
|
||||
target_tuple = await self._aget_tuple_raw(config)
|
||||
cursor_config: RunnableConfig | None = (
|
||||
target_tuple.parent_config if target_tuple else None
|
||||
)
|
||||
while cursor_config is not None:
|
||||
tup = await self._aget_tuple_raw(cursor_config)
|
||||
if tup is None:
|
||||
break
|
||||
if tup.pending_writes:
|
||||
for write in reversed(tup.pending_writes):
|
||||
if write[1] != channel:
|
||||
continue
|
||||
collected.append(write)
|
||||
ancestor_value = tup.checkpoint["channel_values"].get(channel)
|
||||
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=ancestor_value, writes=collected)
|
||||
cursor_config = tup.parent_config
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
|
||||
|
||||
def get_next_version(self, current: V | None, channel: None) -> V:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
|
||||
@@ -14,16 +14,20 @@ from typing import Any
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
PendingWrite,
|
||||
SerializerProtocol,
|
||||
_ChannelWritesHistory,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -121,16 +125,114 @@ class InMemorySaver(
|
||||
return self.stack.__exit__(__exc_type, __exc_value, __traceback)
|
||||
|
||||
def _load_blobs(
|
||||
self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
versions: ChannelVersions,
|
||||
) -> dict[str, Any]:
|
||||
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)
|
||||
return channel_values
|
||||
result: dict[str, Any] = {}
|
||||
for k, ver in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, ver)
|
||||
if kk not in self.blobs:
|
||||
continue
|
||||
vv = self.blobs[kk]
|
||||
if vv[0] == "empty":
|
||||
continue
|
||||
result[k] = self.serde.loads_typed(vv)
|
||||
return result
|
||||
|
||||
def _get_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = config["configurable"].get("checkpoint_id", "")
|
||||
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
|
||||
# Walk the parent chain newest→oldest. Skip the target itself —
|
||||
# writes stored AT `checkpoint_id` are pending for the next step
|
||||
# (pregel applies them via `apply_writes`; they aren't part of the
|
||||
# snapshot value AT `checkpoint_id`).
|
||||
chain: list[str] = []
|
||||
target_entry = ns_storage.get(checkpoint_id)
|
||||
current: str | None = target_entry[2] if target_entry is not None else None
|
||||
while current is not None:
|
||||
entry = ns_storage.get(current)
|
||||
if entry is None:
|
||||
break
|
||||
chain.append(current)
|
||||
_, _, parent = entry
|
||||
current = parent
|
||||
# Scan newest→oldest. A pre-delta blob on an ancestor terminates the
|
||||
# walk and is bound as `seed`; without this, a thread migrated from
|
||||
# pre-delta storage would replay ancestor writes all the way to the
|
||||
# root AND miss any value that lived only in the old blob (e.g. from
|
||||
# `update_state`).
|
||||
#
|
||||
# At each ancestor, check the blob BEFORE processing its pending
|
||||
# writes: a pre-delta blob represents the state AT that ancestor,
|
||||
# which already subsumes any writes stored under it. Processing
|
||||
# those writes first would fold them into the reconstructed value
|
||||
# twice (once via the blob, once via replay).
|
||||
collected: list[PendingWrite] = [] # newest first
|
||||
for cp_id in chain: # newest → oldest
|
||||
entry = ns_storage.get(cp_id)
|
||||
if entry is not None:
|
||||
ckpt = self.serde.loads_typed(entry[0])
|
||||
ver = ckpt.get("channel_versions", {}).get(channel)
|
||||
if ver is not None:
|
||||
blob_entry = self.blobs.get(
|
||||
(thread_id, checkpoint_ns, channel, ver)
|
||||
)
|
||||
if blob_entry is not None and blob_entry[0] != "empty":
|
||||
blob_value = self.serde.loads_typed(blob_entry)
|
||||
if blob_value is not DELTA_SENTINEL:
|
||||
if isinstance(blob_value, _DeltaSnapshot):
|
||||
# Step-based snapshot: the blob is state AT this
|
||||
# ancestor, but the ancestor's pending_writes
|
||||
# encode the NEXT step's transition and are NOT
|
||||
# subsumed by the snapshot — collect them first.
|
||||
step_writes = self.writes.get(
|
||||
(thread_id, checkpoint_ns, cp_id), {}
|
||||
)
|
||||
for (_task_id, _idx), (
|
||||
tid,
|
||||
ch,
|
||||
serialized,
|
||||
_,
|
||||
) in sorted(step_writes.items(), reverse=True):
|
||||
if ch != channel:
|
||||
continue
|
||||
collected.append(
|
||||
(tid, ch, self.serde.loads_typed(serialized))
|
||||
)
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
)
|
||||
# Pre-delta blob: state AT this ancestor already
|
||||
# subsumes its pending_writes — skip them.
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(
|
||||
seed=blob_value, writes=collected
|
||||
)
|
||||
|
||||
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
|
||||
# Within a superstep, sorted by (task_id, idx) = oldest → newest;
|
||||
# reverse for newest-first scan.
|
||||
for (_task_id, _idx), (tid, ch, serialized, _) in sorted(
|
||||
step_writes.items(), reverse=True
|
||||
):
|
||||
if ch != channel:
|
||||
continue
|
||||
val = self.serde.loads_typed(serialized)
|
||||
collected.append((tid, ch, val))
|
||||
collected.reverse()
|
||||
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
|
||||
|
||||
async def _aget_channel_writes_history(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> _ChannelWritesHistory:
|
||||
return self._get_channel_writes_history(config, channel)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the in-memory storage.
|
||||
|
||||
@@ -33,14 +33,18 @@ from langchain_core.load.load import Reviver
|
||||
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
||||
from langgraph.checkpoint.serde.base import SerializerProtocol
|
||||
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
|
||||
from langgraph.checkpoint.serde.types import SendProtocol
|
||||
from langgraph.checkpoint.serde.types import (
|
||||
DELTA_SENTINEL,
|
||||
SendProtocol,
|
||||
_DeltaSentinel,
|
||||
_DeltaSnapshot,
|
||||
)
|
||||
from langgraph.store.base import Item
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.checkpoint.serde._msgpack import (
|
||||
AllowedMsgpackModules,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import SendProtocol
|
||||
|
||||
LC_REVIVER = Reviver()
|
||||
EMPTY_BYTES = b""
|
||||
@@ -55,19 +59,6 @@ _warned_unregistered_types: set[tuple[str, str]] = set()
|
||||
_warned_blocked_types: set[tuple[str, str]] = set()
|
||||
|
||||
|
||||
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:
|
||||
@@ -177,23 +168,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 +228,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}). "
|
||||
@@ -317,10 +297,16 @@ EXT_METHOD_SINGLE_ARG = 3
|
||||
EXT_PYDANTIC_V1 = 4
|
||||
EXT_PYDANTIC_V2 = 5
|
||||
EXT_NUMPY_ARRAY = 6
|
||||
EXT_DELTA_SNAPSHOT = 7
|
||||
EXT_DELTA_SENTINEL = 8
|
||||
|
||||
|
||||
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
||||
if hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
||||
if isinstance(obj, _DeltaSnapshot):
|
||||
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
|
||||
elif isinstance(obj, _DeltaSentinel):
|
||||
return ormsgpack.Ext(EXT_DELTA_SENTINEL, b"")
|
||||
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
||||
return ormsgpack.Ext(
|
||||
EXT_PYDANTIC_V2,
|
||||
_msgpack_enc(
|
||||
@@ -634,7 +620,15 @@ def _create_msgpack_ext_hook(
|
||||
return False
|
||||
|
||||
def ext_hook(code: int, data: bytes) -> Any:
|
||||
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
if code == EXT_DELTA_SENTINEL:
|
||||
return DELTA_SENTINEL
|
||||
elif code == EXT_DELTA_SNAPSHOT:
|
||||
return _DeltaSnapshot(
|
||||
ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
)
|
||||
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from collections.abc import Sequence
|
||||
from typing import (
|
||||
Any,
|
||||
NamedTuple,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
runtime_checkable,
|
||||
@@ -14,6 +15,39 @@ INTERRUPT = "__interrupt__"
|
||||
RESUME = "__resume__"
|
||||
TASKS = "__pregel_tasks"
|
||||
|
||||
|
||||
class _DeltaSentinel:
|
||||
"""Singleton marker stored (as zero bytes) in checkpoint_blobs for a
|
||||
DeltaChannel field. The actual per-step writes live in checkpoint_writes
|
||||
and are replayed through the reducer at load time.
|
||||
|
||||
Compare with `is DELTA_SENTINEL` — `loads_typed` always returns the same
|
||||
module-level instance.
|
||||
"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "DELTA_SENTINEL"
|
||||
|
||||
|
||||
DELTA_SENTINEL = _DeltaSentinel()
|
||||
|
||||
|
||||
class _DeltaSnapshot(NamedTuple):
|
||||
"""Snapshot blob for a DeltaChannel with finite snapshot_frequency.
|
||||
|
||||
Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code.
|
||||
The ancestor walk in `_get_channel_writes_history` terminates when it
|
||||
encounters this type (any non-sentinel blob stops the walk).
|
||||
|
||||
`from_checkpoint` reconstructs the channel value directly from `.value`
|
||||
without replaying writes — the snapshot IS the accumulated state.
|
||||
"""
|
||||
|
||||
value: Any
|
||||
|
||||
|
||||
Value = TypeVar("Value", covariant=True)
|
||||
Update = TypeVar("Update", contravariant=True)
|
||||
C = TypeVar("C")
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.1.0a2"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -333,57 +333,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 = {
|
||||
@@ -1048,3 +997,15 @@ 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_sentinel_serde_round_trip() -> None:
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
type_tag, blob = serde.dumps_typed(DELTA_SENTINEL)
|
||||
assert type_tag == "msgpack"
|
||||
assert blob # non-empty ext envelope
|
||||
loaded = serde.loads_typed((type_tag, blob))
|
||||
assert loaded is DELTA_SENTINEL
|
||||
|
||||
@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
create_checkpoint,
|
||||
@@ -208,8 +209,6 @@ class TestMemorySaver:
|
||||
|
||||
|
||||
async def test_memory_saver() -> None:
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
memory_saver = InMemorySaver()
|
||||
assert isinstance(memory_saver, InMemorySaver)
|
||||
|
||||
@@ -320,3 +319,347 @@ 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_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
|
||||
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
v1 = "00000000000000000000000000000001.0000000000000000"
|
||||
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL)
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp1["channel_versions"][channel] = v1
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
}
|
||||
|
||||
result = saver._load_blobs(thread_id, ns, {channel: v1})
|
||||
assert channel in result
|
||||
assert result[channel] is DELTA_SENTINEL
|
||||
|
||||
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
|
||||
"""_get_channel_writes_history collects ancestor writes oldest→newest,
|
||||
and excludes writes stored at the target checkpoint itself (those are
|
||||
pending writes for the next step, applied separately by pregel)."""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "cp2"
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
||||
}
|
||||
# Writes stored at cp1 produced the cp1 snapshot; part of history.
|
||||
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "hi"}),
|
||||
"",
|
||||
)
|
||||
# Writes stored at cp2 are pending — they will produce cp3 when the
|
||||
# step that loaded cp2 completes. They MUST NOT appear in the
|
||||
# reconstructed snapshot value at cp2.
|
||||
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "pending"}),
|
||||
"",
|
||||
)
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": "cp2",
|
||||
}
|
||||
}
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == [{"content": "hi"}]
|
||||
|
||||
def test_get_channel_writes_at_root_returns_empty(self) -> None:
|
||||
"""Reconstructing the root checkpoint's state: no ancestors → []."""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
}
|
||||
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "pending"}),
|
||||
"",
|
||||
)
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": "cp1",
|
||||
}
|
||||
}
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
assert result.writes == []
|
||||
|
||||
|
||||
class TestBaseFallbackGetChannelWrites:
|
||||
"""Exercises the `BaseCheckpointSaver._get_channel_writes_history` default
|
||||
implementation — the path third-party savers inherit when they don't
|
||||
override `_get_channel_writes_history` themselves.
|
||||
|
||||
Regression guard for a bug where the fallback passed the caller's config
|
||||
(with `checkpoint_id`) straight to `self.list()`, which most savers
|
||||
collapse to a single row — causing the fallback to return `[]`.
|
||||
"""
|
||||
|
||||
def _build_saver_with_chain(self) -> tuple[InMemorySaver, str, str]:
|
||||
"""Build an InMemorySaver with a 3-checkpoint chain and per-step writes
|
||||
for a `messages` channel.
|
||||
|
||||
Returns `(saver, thread_id, namespace)`. The saver subclass deletes the
|
||||
InMemorySaver override so the base class fallback is exercised.
|
||||
"""
|
||||
|
||||
class _ThirdPartyStyleSaver(InMemorySaver):
|
||||
_get_channel_writes_history = (
|
||||
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
_aget_channel_writes_history = (
|
||||
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
saver = _ThirdPartyStyleSaver()
|
||||
serde = JsonPlusSerializer()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
cp0 = empty_checkpoint()
|
||||
cp0["id"] = "00000000000000000000000000000001.0000000000000000"
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "00000000000000000000000000000002.0000000000000000"
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "00000000000000000000000000000003.0000000000000000"
|
||||
saver.storage[thread_id][ns] = {
|
||||
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
|
||||
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
|
||||
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]),
|
||||
}
|
||||
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
|
||||
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "first"}),
|
||||
"",
|
||||
)
|
||||
saver.writes[(thread_id, ns, cp1["id"])][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "second"}),
|
||||
"",
|
||||
)
|
||||
return saver, thread_id, ns
|
||||
|
||||
def test_fallback_returns_ancestor_writes_oldest_first(self) -> None:
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._get_channel_writes_history(config, "messages")
|
||||
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == [{"content": "first"}, {"content": "second"}]
|
||||
|
||||
async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None:
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
result = await saver._aget_channel_writes_history(config, "messages")
|
||||
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == [{"content": "first"}, {"content": "second"}]
|
||||
|
||||
async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None:
|
||||
"""Regression: the re-entrancy guard must be task-local, not thread-local.
|
||||
|
||||
Two concurrent `_aget_channel_writes_history` calls on the same
|
||||
event-loop thread must each see their full reconstructed writes. A
|
||||
`threading.local()` guard would let whichever task set it first
|
||||
short-circuit the other to `writes=[]`.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
saver, thread_id, ns = self._build_saver_with_chain()
|
||||
|
||||
# Force the two tasks to interleave across the `set(True)` boundary:
|
||||
# each `aget_tuple` yields control, so if the guard were thread-local
|
||||
# the second task would observe `active=True` set by the first.
|
||||
orig_aget_tuple = saver.aget_tuple
|
||||
|
||||
async def slow_aget_tuple(config: RunnableConfig) -> Any:
|
||||
await asyncio.sleep(0)
|
||||
return await orig_aget_tuple(config)
|
||||
|
||||
saver.aget_tuple = slow_aget_tuple # type: ignore[method-assign]
|
||||
|
||||
target_id = "00000000000000000000000000000003.0000000000000000"
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target_id,
|
||||
}
|
||||
}
|
||||
|
||||
results = await asyncio.gather(
|
||||
saver._aget_channel_writes_history(config, "messages"),
|
||||
saver._aget_channel_writes_history(config, "messages"),
|
||||
)
|
||||
|
||||
expected_values = [{"content": "first"}, {"content": "second"}]
|
||||
for result in results:
|
||||
assert result.seed is DELTA_SENTINEL
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == expected_values
|
||||
|
||||
|
||||
class TestPreDeltaBlobTerminator:
|
||||
"""Verify the pre-delta blob terminator: when the ancestor walk hits a
|
||||
checkpoint whose blob for the channel is a real value (not
|
||||
DELTA_SENTINEL), reconstruction seeds from it and stops. This guards
|
||||
|
||||
* back-compat: a thread written by pre-delta code, then extended under
|
||||
delta — reconstruction must return the correct value without walking
|
||||
past the last pre-delta ancestor;
|
||||
* perf: without the terminator, every reconstruct-after-migration would
|
||||
walk all the way to the thread root.
|
||||
"""
|
||||
|
||||
def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]:
|
||||
"""Three-checkpoint chain: cp1 (pre-delta, blob=[A]), cp2 (delta,
|
||||
write=B), cp3 (delta, write=C). Reconstructing at cp3 must yield
|
||||
seed=[A] + writes=[B, C].
|
||||
|
||||
Returns `(saver, thread_id, ns, channel, cp3_id)`.
|
||||
"""
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
|
||||
v1 = "00000000000000000000000000000001.0"
|
||||
v2 = "00000000000000000000000000000002.0"
|
||||
v3 = "00000000000000000000000000000003.0"
|
||||
|
||||
# Pre-delta: cp1 stored a real blob for the channel.
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"])
|
||||
# Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes.
|
||||
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL)
|
||||
saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL)
|
||||
|
||||
cp1 = empty_checkpoint()
|
||||
cp1["id"] = "cp1"
|
||||
cp1["channel_versions"][channel] = v1
|
||||
cp2 = empty_checkpoint()
|
||||
cp2["id"] = "cp2"
|
||||
cp2["channel_versions"][channel] = v2
|
||||
cp3 = empty_checkpoint()
|
||||
cp3["id"] = "cp3"
|
||||
cp3["channel_versions"][channel] = v3
|
||||
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
||||
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"),
|
||||
}
|
||||
# Write under cp1 would be from the pre-delta era and MUST be ignored
|
||||
# (the blob already captures it). We add one and assert it is not
|
||||
# folded into the reconstructed result.
|
||||
saver.writes[(thread_id, ns, "cp1")][("task0", 0)] = (
|
||||
"task0",
|
||||
channel,
|
||||
serde.dumps_typed("PRE-DELTA-WRITE"),
|
||||
"",
|
||||
)
|
||||
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed("B"),
|
||||
"",
|
||||
)
|
||||
saver.writes[(thread_id, ns, "cp3")][("task3", 0)] = (
|
||||
"task3",
|
||||
channel,
|
||||
serde.dumps_typed("PENDING-AT-TARGET"),
|
||||
"",
|
||||
)
|
||||
return saver, thread_id, ns, channel, "cp3"
|
||||
|
||||
def test_seed_from_pre_delta_ancestor_blob(self) -> None:
|
||||
saver, thread_id, ns, channel, target = self._build_mixed_thread()
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
|
||||
# Seed came from the pre-delta blob at cp1.
|
||||
assert result.seed == ["A"]
|
||||
# Delta-era writes from cp2 replay through the reducer on top of seed.
|
||||
# cp3 is the target — its own write is pending for the NEXT step and
|
||||
# must be excluded.
|
||||
values = [v for _, _, v in result.writes]
|
||||
assert values == ["B"]
|
||||
|
||||
def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None:
|
||||
"""Writes stored at the pre-delta ancestor itself must not be replayed
|
||||
(the blob subsumes them)."""
|
||||
saver, thread_id, ns, channel, target = self._build_mixed_thread()
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": target,
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._get_channel_writes_history(config, channel)
|
||||
|
||||
values = [v for _, _, v in result.writes]
|
||||
# The pre-delta write under cp1 must not appear (the blob subsumes it).
|
||||
assert "PRE-DELTA-WRITE" not in values
|
||||
# And the pending write at the target is never folded in.
|
||||
assert "PENDING-AT-TARGET" not in values
|
||||
|
||||
Generated
+4
-1
@@ -7,6 +7,9 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
|
||||
[options]
|
||||
prerelease-mode = "allow"
|
||||
|
||||
[[package]]
|
||||
name = "annotated-types"
|
||||
version = "0.7.0"
|
||||
@@ -286,7 +289,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.1.0a2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.4.24"
|
||||
__version__ = "0.4.23"
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -22,10 +22,9 @@ __all__ = ("BinaryOperatorAggregate",)
|
||||
def _strip_extras(t): # type: ignore[no-untyped-def]
|
||||
"""Strips Annotated, Required and NotRequired from a given type."""
|
||||
if hasattr(t, "__origin__"):
|
||||
if t.__origin__ in (Required, NotRequired):
|
||||
return _strip_extras(t.__args__[0])
|
||||
return _strip_extras(t.__origin__)
|
||||
if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
|
||||
return _strip_extras(t.__args__[0])
|
||||
|
||||
return t
|
||||
|
||||
|
||||
@@ -33,11 +32,22 @@ def _get_overwrite(value: Any) -> tuple[bool, Any]:
|
||||
"""Inspects the given value and returns (is_overwrite, overwrite_value)."""
|
||||
if isinstance(value, Overwrite):
|
||||
return True, value.value
|
||||
if isinstance(value, dict) and set(value.keys()) == {OVERWRITE}:
|
||||
if isinstance(value, dict) and len(value) == 1 and OVERWRITE in value:
|
||||
return True, value[OVERWRITE]
|
||||
return False, None
|
||||
|
||||
|
||||
def _operators_equal(a: Callable, b: Callable) -> bool:
|
||||
"""Return True if two reducer operators should be considered equal.
|
||||
|
||||
Lambdas all share the name '<lambda>' so identity comparison is
|
||||
unreliable; treat any pairing that includes a lambda as equal.
|
||||
"""
|
||||
if a.__name__ == "<lambda>" or b.__name__ == "<lambda>":
|
||||
return True
|
||||
return a is b
|
||||
|
||||
|
||||
class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
|
||||
"""Stores the result of applying a binary operator to the current value and each new value.
|
||||
|
||||
@@ -68,11 +78,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
|
||||
self.value = MISSING
|
||||
|
||||
def __eq__(self, value: object) -> bool:
|
||||
return isinstance(value, BinaryOperatorAggregate) and (
|
||||
value.operator is self.operator
|
||||
if value.operator.__name__ != "<lambda>"
|
||||
and self.operator.__name__ != "<lambda>"
|
||||
else True
|
||||
return isinstance(value, BinaryOperatorAggregate) and _operators_equal(
|
||||
self.operator, value.operator
|
||||
)
|
||||
|
||||
@property
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import collections.abc
|
||||
import copy as _copy
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel, Value
|
||||
from langgraph.channels.binop import _get_overwrite, _operators_equal, _strip_extras
|
||||
from langgraph.errors import (
|
||||
EmptyChannelError,
|
||||
ErrorCode,
|
||||
InvalidUpdateError,
|
||||
create_error_message,
|
||||
)
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
"""Reducer channel that stores only a sentinel in checkpoint blobs and
|
||||
reconstructs state by replaying ancestor writes through the reducer.
|
||||
|
||||
The reducer receives the current accumulated value and a batch of writes
|
||||
in one call: `reducer(state, [write1, write2, ...]) -> new_state`.
|
||||
|
||||
Reducers must be deterministic and batching-invariant (associative across
|
||||
folds): applying two consecutive write batches separately must produce the
|
||||
same state as applying their concatenation once:
|
||||
|
||||
reducer(reducer(state, xs), ys) == reducer(state, xs + ys)
|
||||
|
||||
This lets LangGraph replay checkpointed writes in larger batches than they
|
||||
were originally produced without changing reconstructed state.
|
||||
|
||||
`snapshot_frequency=None` (default): pure delta; stores only
|
||||
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes.
|
||||
|
||||
`snapshot_frequency=N`: `create_checkpoint` writes a full `_DeltaSnapshot`
|
||||
blob every N steps, bounding replay depth to N.
|
||||
|
||||
Parameters:
|
||||
reducer: `(state, list[writes]) -> new_state`. Must be deterministic
|
||||
and batching-invariant as described above.
|
||||
typ: The value type (e.g. `list`, `dict`). Inferred automatically
|
||||
from the outer type when used inside `Annotated[T, DeltaChannel(...)]`.
|
||||
snapshot_frequency: Every Nth pregel step writes a snapshot blob.
|
||||
`None` (default) = pure delta, never snapshot.
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "reducer", "snapshot_frequency")
|
||||
value: Value | Any
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reducer: Callable[[Any, Sequence[Any]], Any],
|
||||
typ: type[Value] | None = None,
|
||||
*,
|
||||
snapshot_frequency: int | None = None,
|
||||
) -> None:
|
||||
if typ is None:
|
||||
typ = list # type: ignore[assignment] # placeholder; overridden by _is_field_channel
|
||||
super().__init__(typ)
|
||||
self.reducer = reducer
|
||||
self.snapshot_frequency = snapshot_frequency
|
||||
typ = _strip_extras(typ)
|
||||
if typ in (collections.abc.Sequence, collections.abc.MutableSequence):
|
||||
typ = list
|
||||
if typ in (collections.abc.Set, collections.abc.MutableSet):
|
||||
typ = set
|
||||
if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
|
||||
typ = dict
|
||||
self.typ = typ
|
||||
self.value: Any = MISSING
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
return False
|
||||
if self.snapshot_frequency != other.snapshot_frequency:
|
||||
return False
|
||||
return _operators_equal(self.reducer, other.reducer)
|
||||
|
||||
@property
|
||||
def ValueType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ
|
||||
|
||||
def is_snapshot_step(self, step: int) -> bool:
|
||||
"""True if pregel should write a snapshot blob at this step."""
|
||||
return (
|
||||
self.snapshot_frequency is not None
|
||||
and step > 0
|
||||
and step % self.snapshot_frequency == 0
|
||||
)
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = self.__class__(
|
||||
self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
|
||||
)
|
||||
new.key = self.key
|
||||
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
||||
return new
|
||||
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
"""Initialize from a stored blob or sentinel.
|
||||
|
||||
Blob types (dispatched via serde ext code, not dict key inspection):
|
||||
* `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes.
|
||||
* `_DeltaSnapshot(value)`: restore value directly from snapshot.
|
||||
* plain value (migration from old BinOp blobs): use directly.
|
||||
"""
|
||||
new = self.__class__(
|
||||
self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
|
||||
)
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
|
||||
new.value = self.typ()
|
||||
elif isinstance(checkpoint, _DeltaSnapshot):
|
||||
new.value = checkpoint.value
|
||||
else:
|
||||
new.value = checkpoint
|
||||
return new
|
||||
|
||||
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
|
||||
"""Apply ancestor writes oldest-to-newest via a single reducer call.
|
||||
|
||||
If any write is an Overwrite, the last one in the sequence acts as
|
||||
the reset point: its value becomes the new base and only writes
|
||||
after it are passed to the reducer.
|
||||
"""
|
||||
values = [v for _, _, v in writes]
|
||||
if not values:
|
||||
return
|
||||
base = self.value
|
||||
start = 0
|
||||
for i, v in enumerate(values):
|
||||
is_ow, ow_value = _get_overwrite(v)
|
||||
if is_ow:
|
||||
base = _copy.copy(ow_value) if ow_value is not None else self.typ()
|
||||
start = i + 1
|
||||
remaining = values[start:]
|
||||
self.value = self.reducer(base, remaining) if remaining else base
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
if not values:
|
||||
return False
|
||||
overwrite_idx: int | None = None
|
||||
for i, v in enumerate(values):
|
||||
is_ow, _ = _get_overwrite(v)
|
||||
if is_ow:
|
||||
if overwrite_idx is not None:
|
||||
msg = create_error_message(
|
||||
message="Can receive only one Overwrite value per super-step.",
|
||||
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
|
||||
)
|
||||
raise InvalidUpdateError(msg)
|
||||
overwrite_idx = i
|
||||
if overwrite_idx is not None:
|
||||
_, overwrite_value = _get_overwrite(values[overwrite_idx])
|
||||
base = (
|
||||
_copy.copy(overwrite_value)
|
||||
if overwrite_value is not None
|
||||
else self.typ()
|
||||
)
|
||||
remaining = [v for i, v in enumerate(values) if i != overwrite_idx]
|
||||
self.value = self.reducer(base, remaining) if remaining else base
|
||||
return True
|
||||
base = self.typ() if self.value is MISSING else self.value
|
||||
self.value = self.reducer(base, list(values))
|
||||
return True
|
||||
|
||||
def get(self) -> Any:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
"""Return stored representation: always `DELTA_SENTINEL`.
|
||||
|
||||
Snapshot decisions are made by `create_checkpoint` in pregel (which
|
||||
has the step number) via `is_snapshot_step`. `checkpoint()` is only
|
||||
called for non-snapshot steps or when no checkpointer is available.
|
||||
"""
|
||||
if self.value is MISSING:
|
||||
return MISSING
|
||||
return DELTA_SENTINEL
|
||||
@@ -244,6 +244,52 @@ def add_messages(
|
||||
return merged
|
||||
|
||||
|
||||
def _messages_delta_reducer(
|
||||
state: list[AnyMessage], writes: list[list[AnyMessage]]
|
||||
) -> list[AnyMessage]:
|
||||
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
|
||||
|
||||
Processes all writes in one pass — dedup by ID, `RemoveMessage`
|
||||
tombstoning — without calling `add_messages`. Assumes writes contain
|
||||
already-typed `BaseMessage` objects (no raw-dict coercion).
|
||||
|
||||
This reducer is batching-invariant, as required by `DeltaChannel`:
|
||||
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
|
||||
|
||||
Use `add_messages` as the reducer for `BinaryOperatorAggregate` or
|
||||
anywhere raw message dicts / strings need to be coerced first.
|
||||
|
||||
Example::
|
||||
|
||||
from typing import Annotated
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
"""
|
||||
from itertools import chain
|
||||
|
||||
index: dict[str, int] = {m.id: i for i, m in enumerate(state) if m.id is not None}
|
||||
result: list[AnyMessage | None] = list(state)
|
||||
for msg in chain.from_iterable(
|
||||
[w] if isinstance(w, BaseMessage) else w for w in writes
|
||||
):
|
||||
mid = msg.id
|
||||
if mid is None:
|
||||
result.append(msg)
|
||||
elif isinstance(msg, RemoveMessage):
|
||||
if mid in index:
|
||||
result[index[mid]] = None
|
||||
del index[mid]
|
||||
elif mid in index:
|
||||
result[index[mid]] = msg
|
||||
else:
|
||||
index[mid] = len(result)
|
||||
result.append(msg)
|
||||
return [m for m in result if m is not None]
|
||||
|
||||
|
||||
@deprecated(
|
||||
"MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",
|
||||
category=None,
|
||||
|
||||
@@ -48,6 +48,7 @@ from langgraph._internal._runnable import coerce_to_runnable
|
||||
from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
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 (
|
||||
@@ -1082,6 +1083,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
CompiledStateGraph: The compiled `StateGraph`.
|
||||
"""
|
||||
checkpointer = ensure_valid_checkpointer(checkpointer)
|
||||
|
||||
serde_allowlist: set[tuple[str, ...]] | None = None
|
||||
if _serde.STRICT_MSGPACK_ENABLED:
|
||||
schema_types: list[type[Any]] = [
|
||||
@@ -1667,6 +1669,20 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None:
|
||||
# Search through all annotated medata to find channel annotations
|
||||
for item in meta:
|
||||
if isinstance(item, BaseChannel):
|
||||
if isinstance(item, DeltaChannel) and hasattr(typ, "__origin__"):
|
||||
origin = typ.__origin__
|
||||
# Unwrap parameterized Required[X]/NotRequired[X] to X
|
||||
# (e.g. Annotated[NotRequired[dict[...]], ...]).
|
||||
if hasattr(origin, "__origin__") and origin.__origin__ in (
|
||||
Required,
|
||||
NotRequired,
|
||||
):
|
||||
origin = origin.__args__[0]
|
||||
item = item.__class__(
|
||||
item.reducer,
|
||||
origin,
|
||||
snapshot_frequency=item.snapshot_frequency,
|
||||
)
|
||||
return item
|
||||
elif isclass(item) and issubclass(item, BaseChannel):
|
||||
# ex, Annotated[int, EphemeralValue, SomeOtherAnnotation]
|
||||
|
||||
@@ -1,17 +1,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, cast
|
||||
|
||||
from langgraph.checkpoint.base import Checkpoint
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
LATEST_VERSION = 4
|
||||
|
||||
GetNextVersion = Callable[[Any, None], Any]
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
@@ -31,35 +37,87 @@ def create_checkpoint(
|
||||
*,
|
||||
id: str | None = None,
|
||||
updated_channels: set[str] | None = None,
|
||||
get_next_version: GetNextVersion | None = None,
|
||||
force_delta_snapshot: bool = False,
|
||||
) -> Checkpoint:
|
||||
"""Create a checkpoint for the given channels."""
|
||||
"""Create a checkpoint for the given channels.
|
||||
|
||||
For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a
|
||||
`_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor
|
||||
walk to at most N steps. Snapshots are eager: even if the channel had no
|
||||
write this step, a version bump is forced (via `get_next_version`) so the
|
||||
blob is stored by `put()`. Without `get_next_version` (e.g. static
|
||||
contexts), snapshot steps gracefully fall back to sentinel.
|
||||
|
||||
`force_delta_snapshot` writes available `DeltaChannel` values as snapshots
|
||||
regardless of `snapshot_frequency`. This is used by `durability="exit"`,
|
||||
where intermediate writes are not stored as ancestor `checkpoint_writes`.
|
||||
"""
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
if channels is None:
|
||||
values = checkpoint["channel_values"]
|
||||
channel_versions = checkpoint["channel_versions"]
|
||||
else:
|
||||
values = {}
|
||||
channel_versions = dict(checkpoint["channel_versions"])
|
||||
for k in channels:
|
||||
if k not in checkpoint["channel_versions"]:
|
||||
if k not in channel_versions:
|
||||
continue
|
||||
v = channels[k].checkpoint()
|
||||
if v is not MISSING:
|
||||
values[k] = v
|
||||
ch = channels[k]
|
||||
if (
|
||||
isinstance(ch, DeltaChannel)
|
||||
and (force_delta_snapshot or ch.is_snapshot_step(step))
|
||||
and ch.is_available()
|
||||
):
|
||||
# Eager snapshot: bump version if not already written this step
|
||||
# so put() includes this channel in new_versions and stores blob.
|
||||
if get_next_version is not None and (
|
||||
updated_channels is None or k not in updated_channels
|
||||
):
|
||||
channel_versions[k] = get_next_version(channel_versions[k], None)
|
||||
values[k] = _DeltaSnapshot(ch.get())
|
||||
else:
|
||||
v = ch.checkpoint()
|
||||
if v is not MISSING:
|
||||
values[k] = v
|
||||
return Checkpoint(
|
||||
v=LATEST_VERSION,
|
||||
ts=ts,
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
channel_versions=channel_versions,
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
updated_channels=None if updated_channels is None else sorted(updated_channels),
|
||||
)
|
||||
|
||||
|
||||
def _needs_replay(spec: BaseChannel, stored: object) -> bool:
|
||||
"""True if `spec` is a `DeltaChannel` and the stored blob is a sentinel,
|
||||
requiring an ancestor walk to reconstruct.
|
||||
|
||||
`_DeltaSnapshot` blobs and plain values (migration) resolve directly via
|
||||
`from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay.
|
||||
"""
|
||||
if not isinstance(spec, DeltaChannel):
|
||||
return False
|
||||
return stored is MISSING or stored is DELTA_SENTINEL
|
||||
|
||||
|
||||
def channels_from_checkpoint(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
checkpoint: Checkpoint,
|
||||
*,
|
||||
saver: BaseCheckpointSaver | None = None,
|
||||
config: RunnableConfig | None = None,
|
||||
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
|
||||
"""Get channels from a checkpoint."""
|
||||
"""Hydrate channels from a checkpoint.
|
||||
|
||||
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
|
||||
is sufficient. `DeltaChannel` is the exception: sentinel blobs require an
|
||||
ancestor walk via `saver._get_channel_writes_history`. The walk terminates
|
||||
at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain
|
||||
value, so read depth is bounded by `snapshot_frequency`.
|
||||
"""
|
||||
channel_specs: dict[str, BaseChannel] = {}
|
||||
managed_specs: dict[str, ManagedValueSpec] = {}
|
||||
for k, v in specs.items():
|
||||
@@ -67,13 +125,53 @@ 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, spec in channel_specs.items():
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
if _needs_replay(spec, stored) and saver is not None and config is not None:
|
||||
delta_spec = cast(DeltaChannel, spec)
|
||||
history = saver._get_channel_writes_history(config, k)
|
||||
replay_ch = delta_spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
ch = replay_ch
|
||||
else:
|
||||
ch = spec.from_checkpoint(stored)
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
|
||||
|
||||
async def achannels_from_checkpoint(
|
||||
specs: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
checkpoint: Checkpoint,
|
||||
*,
|
||||
saver: BaseCheckpointSaver | None = None,
|
||||
config: RunnableConfig | None = None,
|
||||
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
|
||||
"""Async version of `channels_from_checkpoint`. See docstring there."""
|
||||
channel_specs: dict[str, BaseChannel] = {}
|
||||
managed_specs: dict[str, ManagedValueSpec] = {}
|
||||
for k, v in specs.items():
|
||||
if isinstance(v, BaseChannel):
|
||||
channel_specs[k] = v
|
||||
else:
|
||||
managed_specs[k] = v
|
||||
|
||||
channels: dict[str, BaseChannel] = {}
|
||||
for k, spec in channel_specs.items():
|
||||
ch: BaseChannel
|
||||
stored = checkpoint["channel_values"].get(k, MISSING)
|
||||
if _needs_replay(spec, stored) and saver is not None and config is not None:
|
||||
delta_spec = cast(DeltaChannel, spec)
|
||||
history = await saver._aget_channel_writes_history(config, k)
|
||||
replay_ch = delta_spec.from_checkpoint(history.seed)
|
||||
replay_ch.replay_writes(history.writes)
|
||||
ch = replay_ch
|
||||
else:
|
||||
ch = spec.from_checkpoint(stored)
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
|
||||
|
||||
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
|
||||
@@ -68,6 +68,7 @@ from langgraph.callbacks import (
|
||||
GraphResumeEvent,
|
||||
)
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.untracked_value import UntrackedValue
|
||||
from langgraph.constants import TAG_HIDDEN
|
||||
from langgraph.errors import (
|
||||
@@ -92,6 +93,7 @@ from langgraph.pregel._algo import (
|
||||
task_path_str,
|
||||
)
|
||||
from langgraph.pregel._checkpoint import (
|
||||
achannels_from_checkpoint,
|
||||
channels_from_checkpoint,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
@@ -188,6 +190,8 @@ class PregelLoop:
|
||||
_migrate_checkpoint: Callable[[Checkpoint], None] | None
|
||||
submit: Submit
|
||||
channels: Mapping[str, BaseChannel]
|
||||
# Only set on AsyncPregelLoop; sync loops keep this as None.
|
||||
_delta_write_futs: list[Any] | None = None
|
||||
managed: ManagedValueMapping
|
||||
checkpoint: Checkpoint
|
||||
checkpoint_id_saved: str
|
||||
@@ -406,7 +410,7 @@ class PregelLoop:
|
||||
task = self.tasks.get(task_id)
|
||||
else:
|
||||
task = None
|
||||
self.submit(
|
||||
fut = self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
config,
|
||||
writes_to_save,
|
||||
@@ -414,12 +418,16 @@ class PregelLoop:
|
||||
task_path_str(task.path) if task else "",
|
||||
)
|
||||
else:
|
||||
self.submit(
|
||||
fut = self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
config,
|
||||
writes_to_save,
|
||||
task_id,
|
||||
)
|
||||
if self._delta_write_futs is not None and any(
|
||||
isinstance(self.specs.get(c), DeltaChannel) for c, _ in writes_to_save
|
||||
):
|
||||
self._delta_write_futs.append(fut)
|
||||
# output writes
|
||||
if hasattr(self, "tasks"):
|
||||
self.output_writes(task_id, writes)
|
||||
@@ -890,6 +898,10 @@ class PregelLoop:
|
||||
self.step,
|
||||
id=self.checkpoint["id"] if exiting else None,
|
||||
updated_channels=self.updated_channels,
|
||||
get_next_version=self.checkpointer_get_next_version
|
||||
if do_checkpoint
|
||||
else None,
|
||||
force_delta_snapshot=exiting and self.durability == "exit",
|
||||
)
|
||||
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
|
||||
if TASKS in self.checkpoint["channel_values"] and any(
|
||||
@@ -1273,7 +1285,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
)
|
||||
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs, self.checkpoint
|
||||
self.specs,
|
||||
self.checkpoint,
|
||||
saver=self.checkpointer,
|
||||
config=self.checkpoint_config,
|
||||
)
|
||||
self.stack.push(self._suppress_interrupt)
|
||||
self.status = "input"
|
||||
@@ -1368,6 +1383,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
# Drain DeltaChannel write futures before committing the checkpoint so
|
||||
# DELTA_SENTINEL blobs are never saved ahead of their backing writes.
|
||||
if self._delta_write_futs:
|
||||
futs, self._delta_write_futs = self._delta_write_futs, []
|
||||
await asyncio.gather(*futs)
|
||||
try:
|
||||
if prev is not None:
|
||||
await prev
|
||||
@@ -1473,11 +1493,15 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
if saved.pending_writes is not None
|
||||
else []
|
||||
)
|
||||
self._delta_write_futs = []
|
||||
self.submit = await self.stack.enter_async_context(
|
||||
AsyncBackgroundExecutor(self.config)
|
||||
)
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs, self.checkpoint
|
||||
self.channels, self.managed = await achannels_from_checkpoint(
|
||||
self.specs,
|
||||
self.checkpoint,
|
||||
saver=self.checkpointer,
|
||||
config=self.checkpoint_config,
|
||||
)
|
||||
self.stack.push(self._suppress_interrupt)
|
||||
self.status = "input"
|
||||
|
||||
@@ -122,6 +122,7 @@ from langgraph.pregel._algo import (
|
||||
)
|
||||
from langgraph.pregel._call import identifier
|
||||
from langgraph.pregel._checkpoint import (
|
||||
achannels_from_checkpoint,
|
||||
channels_from_checkpoint,
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
@@ -1052,6 +1053,10 @@ class Pregel(
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
saver=self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
@@ -1168,9 +1173,13 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
channels, managed = channels_from_checkpoint(
|
||||
channels, managed = await achannels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
saver=self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
@@ -1541,6 +1550,11 @@ class Pregel(
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saver=self.checkpointer
|
||||
if saved is not None
|
||||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config if saved is not None else None,
|
||||
)
|
||||
values, as_node = updates[0][:2]
|
||||
|
||||
@@ -1984,9 +1998,14 @@ class Pregel(
|
||||
)
|
||||
if saved:
|
||||
checkpoint_config = patch_configurable(config, saved.config[CONF])
|
||||
channels, managed = channels_from_checkpoint(
|
||||
channels, managed = await achannels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saver=self.checkpointer
|
||||
if saved is not None
|
||||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
config=saved.config if saved is not None else None,
|
||||
)
|
||||
values, as_node = updates[0][:2]
|
||||
# no values, just clear all tasks
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.2.0a2"
|
||||
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",
|
||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||
"langchain-core>=1.3.2,<2",
|
||||
"langgraph-checkpoint>=4.0.3,<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",
|
||||
]
|
||||
|
||||
@@ -1,18 +1,34 @@
|
||||
import operator
|
||||
from collections.abc import Sequence
|
||||
from typing import Annotated
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.topic import Topic
|
||||
from langgraph.channels.untracked_value import UntrackedValue
|
||||
from langgraph.errors import EmptyChannelError, InvalidUpdateError
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.graph.state import _get_channel
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core channel primitives
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_last_value() -> None:
|
||||
channel = LastValue(int).from_checkpoint(MISSING)
|
||||
assert channel.ValueType is int
|
||||
@@ -95,25 +111,543 @@ def test_untracked_value() -> None:
|
||||
assert channel.ValueType is dict
|
||||
assert channel.UpdateType is dict
|
||||
|
||||
# UntrackedValue should start empty
|
||||
with pytest.raises(EmptyChannelError):
|
||||
channel.get()
|
||||
|
||||
# Should be able to update with a value
|
||||
test_data = {"session": "test", "temp": "dir"}
|
||||
channel.update([test_data])
|
||||
assert channel.get() == test_data
|
||||
|
||||
# Update with new value
|
||||
new_data = {"session": "updated", "temp": "newdir"}
|
||||
channel.update([new_data])
|
||||
assert channel.get() == new_data
|
||||
|
||||
# On checkpoint, UntrackedValue should return MISSING
|
||||
checkpoint = channel.checkpoint()
|
||||
assert checkpoint is MISSING
|
||||
|
||||
# Creating from checkpoint with MISSING should start empty
|
||||
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
|
||||
with pytest.raises(EmptyChannelError):
|
||||
new_channel.get()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DeltaChannel — message reducer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_channel_basic_two_steps() -> None:
|
||||
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
||||
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert d1 is DELTA_SENTINEL
|
||||
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert d2 is DELTA_SENTINEL
|
||||
|
||||
assert len(ch.get()) == 2
|
||||
assert ch.get()[0].content == "hi"
|
||||
assert ch.get()[1].content == "hello"
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_writes_list() -> None:
|
||||
"""replay_writes on a fresh channel replays through the operator."""
|
||||
spec = DeltaChannel(_messages_delta_reducer, list)
|
||||
ch = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch.replay_writes(
|
||||
[
|
||||
("t0", "messages", HumanMessage(content="hi", id="h1")),
|
||||
("t1", "messages", AIMessage(content="hello", id="a1")),
|
||||
("t2", "messages", HumanMessage(content="bye", id="h2")),
|
||||
]
|
||||
)
|
||||
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:
|
||||
spec = DeltaChannel(_messages_delta_reducer, list)
|
||||
old_value = [HumanMessage(content="old", id="h1")]
|
||||
ch = spec.from_checkpoint(old_value)
|
||||
assert ch.get() == old_value
|
||||
|
||||
|
||||
def test_delta_channel_overwrite() -> None:
|
||||
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
||||
ch.update([HumanMessage(content="old", id="h1")])
|
||||
|
||||
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
|
||||
d = ch.checkpoint()
|
||||
assert d is DELTA_SENTINEL
|
||||
assert len(ch.get()) == 1
|
||||
assert ch.get()[0].content == "new"
|
||||
|
||||
|
||||
def test_delta_channel_remove_message_and_replay() -> None:
|
||||
"""RemoveMessage must round-trip correctly when writes are replayed."""
|
||||
spec = DeltaChannel(_messages_delta_reducer, list)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
assert ch.get() == [
|
||||
HumanMessage(content="hi", id="h1"),
|
||||
AIMessage(content="hello", id="a1"),
|
||||
]
|
||||
|
||||
ch.update([RemoveMessage(id="a1")])
|
||||
assert ch.get() == [HumanMessage(content="hi", id="h1")]
|
||||
|
||||
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch2.replay_writes(
|
||||
[
|
||||
("t0", "messages", HumanMessage(content="hi", id="h1")),
|
||||
("t1", "messages", AIMessage(content="hello", id="a1")),
|
||||
("t2", "messages", RemoveMessage(id="a1")),
|
||||
]
|
||||
)
|
||||
assert ch2.get() == [HumanMessage(content="hi", id="h1")]
|
||||
|
||||
|
||||
def test_delta_channel_update_by_id_and_replay() -> None:
|
||||
"""Updating a message by ID must round-trip correctly through writes replay."""
|
||||
spec = DeltaChannel(_messages_delta_reducer, list)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
|
||||
ch.update([HumanMessage(content="original", id="h1")])
|
||||
ch.update([HumanMessage(content="updated", id="h1")])
|
||||
assert ch.get() == [HumanMessage(content="updated", id="h1")]
|
||||
|
||||
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch2.replay_writes(
|
||||
[
|
||||
("t0", "messages", HumanMessage(content="original", id="h1")),
|
||||
("t1", "messages", HumanMessage(content="updated", id="h1")),
|
||||
]
|
||||
)
|
||||
assert len(ch2.get()) == 1
|
||||
assert ch2.get()[0].content == "updated"
|
||||
|
||||
|
||||
def test_delta_channel_checkpoint_returns_sentinel() -> None:
|
||||
"""checkpoint() always returns DELTA_SENTINEL regardless of state."""
|
||||
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
||||
assert ch.checkpoint() is DELTA_SENTINEL
|
||||
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
assert ch.checkpoint() is DELTA_SENTINEL
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DeltaChannel — snapshot frequency
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_step_based() -> None:
|
||||
"""Snapshots fire on every Nth step regardless of whether the channel was written.
|
||||
|
||||
With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot
|
||||
blob — even if the channel had no write that step (eager snapshot). This
|
||||
bounds the ancestor walk to at most N steps on any read.
|
||||
"""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[
|
||||
list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=5)
|
||||
]
|
||||
other: str
|
||||
|
||||
def node_a(state: State) -> dict:
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
||||
|
||||
def node_b(state: State) -> dict:
|
||||
return {"other": "y"}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("a", node_a)
|
||||
g.add_node("b", node_b)
|
||||
g.add_edge(START, "a")
|
||||
g.add_edge("a", "b")
|
||||
saver = InMemorySaver()
|
||||
graph = g.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
for i in range(6):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""},
|
||||
config,
|
||||
)
|
||||
|
||||
msg_blob_values = [
|
||||
saver.serde.loads_typed((type_tag, blob))
|
||||
for k, (type_tag, blob) in saver.blobs.items()
|
||||
if k[2] == "messages" and type_tag == "msgpack" and blob
|
||||
]
|
||||
snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)]
|
||||
assert snapshots, "expected at least one _DeltaSnapshot blob for messages"
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 12 # 6 human + 6 AI
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_fires_even_when_not_written() -> None:
|
||||
"""Eager snapshot: _DeltaSnapshot stored at snapshot step even when the
|
||||
channel had no write that step (node_b doesn't touch messages).
|
||||
"""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[
|
||||
list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3)
|
||||
]
|
||||
tick: int
|
||||
|
||||
def writer(state: State) -> dict:
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
||||
|
||||
def ticker(state: State) -> dict:
|
||||
return {"tick": state["tick"] + 1}
|
||||
|
||||
g = StateGraph(State)
|
||||
g.add_node("writer", writer)
|
||||
g.add_node("ticker", ticker)
|
||||
g.add_edge(START, "writer")
|
||||
g.add_edge("writer", "ticker")
|
||||
saver = InMemorySaver()
|
||||
graph = g.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
for i in range(5):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0},
|
||||
config,
|
||||
)
|
||||
|
||||
msg_blobs = {
|
||||
k: saver.serde.loads_typed((t, b))
|
||||
for k, (t, b) in saver.blobs.items()
|
||||
if k[2] == "messages" and t == "msgpack" and b
|
||||
}
|
||||
snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)}
|
||||
assert snapshots, (
|
||||
"eager snapshots must fire even on steps where messages wasn't written"
|
||||
)
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 10 # 5 human + 5 AI
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DeltaChannel — end-to-end (InMemorySaver)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer, list)]
|
||||
|
||||
n = {"v": 0}
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n["v"] += 1
|
||||
return {"messages": [AIMessage(content=f"ok{n['v']}", id=f"ai{n['v']}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "t1"}}
|
||||
|
||||
graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
|
||||
graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config)
|
||||
|
||||
saved = saver.get_tuple(config)
|
||||
assert saved is not None
|
||||
assert "messages" in saved.checkpoint["channel_values"]
|
||||
assert saved.checkpoint["channel_values"]["messages"] is DELTA_SENTINEL
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DeltaChannel — dict reducer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _delta_channel_with_type(op, typ):
|
||||
"""Build a DeltaChannel with an explicit type via the Annotated injection path."""
|
||||
return _get_channel("_test", Annotated[typ, DeltaChannel(op)])
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_fresh_channel() -> None:
|
||||
"""DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
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."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
|
||||
ch.update([{"a": 1}])
|
||||
d1 = ch.checkpoint()
|
||||
assert d1 is DELTA_SENTINEL
|
||||
|
||||
ch.update([{"b": 2}])
|
||||
d2 = ch.checkpoint()
|
||||
assert d2 is DELTA_SENTINEL
|
||||
|
||||
assert ch.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_writes_reconstruction() -> None:
|
||||
"""replay_writes on a fresh channel replays through a dict merge reducer."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
spec = _delta_channel_with_type(merge_dicts, dict)
|
||||
ch = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch.replay_writes(
|
||||
[
|
||||
("t0", "files", {"a": 1}),
|
||||
("t1", "files", {"b": 2}),
|
||||
("t2", "files", {"c": 3}),
|
||||
]
|
||||
)
|
||||
assert ch.get() == {"a": 1, "b": 2, "c": 3}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_deletions() -> None:
|
||||
"""Dict reducer that treats None values as deletions works end-to-end."""
|
||||
|
||||
def merge_files(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
for k, v in w.items():
|
||||
if v is None:
|
||||
result.pop(k, None)
|
||||
else:
|
||||
result[k] = v
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING)
|
||||
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
|
||||
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
||||
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
spec = _delta_channel_with_type(merge_files, dict)
|
||||
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch2.replay_writes(
|
||||
[
|
||||
("t0", "files", {"file1.py": "content1", "file2.py": "content2"}),
|
||||
("t1", "files", {"file1.py": None, "file3.py": "content3"}),
|
||||
]
|
||||
)
|
||||
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_overwrite_in_update() -> None:
|
||||
"""Overwrite(dict) in update() must preserve dict shape, not coerce to list."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
ch.update([{"a": 1}])
|
||||
ch.update([Overwrite({"b": 2, "c": 3})])
|
||||
assert ch.get() == {"b": 2, "c": 3}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
|
||||
"""Overwrite(dict) embedded in replayed writes must reconstruct as dict."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
spec = _delta_channel_with_type(merge_dicts, dict)
|
||||
ch = spec.from_checkpoint(DELTA_SENTINEL)
|
||||
ch.replay_writes(
|
||||
[
|
||||
("t0", "files", {"a": 1}),
|
||||
("t1", "files", Overwrite({"x": 10, "y": 20})),
|
||||
("t2", "files", {"z": 30}),
|
||||
]
|
||||
)
|
||||
assert ch.get() == {"x": 10, "y": 20, "z": 30}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
|
||||
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
annotation = Annotated[NotRequired[dict[str, int]], DeltaChannel(merge_dicts)]
|
||||
ch = _get_channel("files", annotation).from_checkpoint(MISSING)
|
||||
assert ch.get() == {}
|
||||
ch.update([{"a": 1}])
|
||||
ch.update([{"b": 2}])
|
||||
assert ch.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
|
||||
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel."""
|
||||
|
||||
def merge_files(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
for k, v in w.items():
|
||||
if v is None:
|
||||
result.pop(k, None)
|
||||
else:
|
||||
result[k] = v
|
||||
return result
|
||||
|
||||
class State(TypedDict):
|
||||
files: Annotated[dict[str, str], DeltaChannel(merge_files)]
|
||||
|
||||
turn = {"v": 0}
|
||||
|
||||
def write_file(state: State) -> dict:
|
||||
turn["v"] += 1
|
||||
n = turn["v"]
|
||||
return {"files": {f"/doc_{n}.txt": f"content for turn {n}"}}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("write_file", write_file)
|
||||
builder.add_edge(START, "write_file")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "fs"}}
|
||||
|
||||
for _ in range(3):
|
||||
graph.invoke({"files": {}}, config)
|
||||
|
||||
saved = saver.get_tuple(config)
|
||||
assert saved is not None
|
||||
assert saved.checkpoint["channel_values"]["files"] is DELTA_SENTINEL
|
||||
state = graph.get_state(config)
|
||||
assert state.values["files"] == {
|
||||
"/doc_1.txt": "content for turn 1",
|
||||
"/doc_2.txt": "content for turn 2",
|
||||
"/doc_3.txt": "content for turn 3",
|
||||
}
|
||||
|
||||
def delete_file(state: State) -> dict:
|
||||
return {"files": {"/doc_1.txt": None}}
|
||||
|
||||
builder2 = StateGraph(State)
|
||||
builder2.add_node("write_file", write_file)
|
||||
builder2.add_node("delete_file", delete_file)
|
||||
builder2.add_edge(START, "write_file")
|
||||
builder2.add_edge("write_file", "delete_file")
|
||||
turn["v"] = 0
|
||||
saver2 = InMemorySaver()
|
||||
graph2 = builder2.compile(checkpointer=saver2)
|
||||
config2 = {"configurable": {"thread_id": "fs2"}}
|
||||
graph2.invoke({"files": {}}, config2)
|
||||
state2 = graph2.get_state(config2)
|
||||
assert state2.values["files"] == {}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_backwards_compat() -> None:
|
||||
"""A pre-DeltaChannel dict checkpoint must load as a dict, not be listified."""
|
||||
|
||||
def merge_dicts(state: dict, writes: list) -> dict:
|
||||
result = dict(state)
|
||||
for w in writes:
|
||||
result.update(w)
|
||||
return result
|
||||
|
||||
spec = _delta_channel_with_type(merge_dicts, dict)
|
||||
old_value = {"a": 1, "b": 2}
|
||||
ch = spec.from_checkpoint(old_value)
|
||||
assert ch.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DeltaChannel — seed / pre-delta migration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_honors_seed() -> None:
|
||||
"""A non-sentinel value to from_checkpoint is used as the pre-delta seed.
|
||||
|
||||
Guards the pre-delta migration path: when the saver's ancestor walk hits
|
||||
a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs
|
||||
the post-migration state correctly rather than replaying from empty.
|
||||
"""
|
||||
spec = DeltaChannel(_messages_delta_reducer, list)
|
||||
seed = [HumanMessage(content="pre-delta", id="p1")]
|
||||
ch = spec.from_checkpoint(seed)
|
||||
ch.replay_writes(
|
||||
[
|
||||
("t0", "messages", AIMessage(content="delta-1", id="d1")),
|
||||
("t1", "messages", HumanMessage(content="delta-2", id="d2")),
|
||||
]
|
||||
)
|
||||
msgs = ch.get()
|
||||
assert [m.content for m in msgs] == ["pre-delta", "delta-1", "delta-2"]
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_seed_without_writes() -> None:
|
||||
"""Reconstruction at a pre-delta ancestor with no newer deltas returns
|
||||
just the seed — the saver's terminator fired immediately."""
|
||||
spec = DeltaChannel(_messages_delta_reducer, list)
|
||||
seed = [HumanMessage(content="only-snap", id="s1")]
|
||||
ch = spec.from_checkpoint(seed)
|
||||
ch.replay_writes([])
|
||||
assert ch.get() == seed
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_sentinel() -> None:
|
||||
"""`seed=None` must start replay from None, not from an empty channel.
|
||||
|
||||
The DELTA_SENTINEL / MISSING sentinels mean 'no seed'; passing `None`
|
||||
explicitly should feed None to the reducer as the left operand.
|
||||
"""
|
||||
|
||||
def replace(state, writes):
|
||||
return writes[-1] if writes else state
|
||||
|
||||
spec = DeltaChannel(replace, list)
|
||||
ch = spec.from_checkpoint(None)
|
||||
ch.replay_writes([("t0", "x", "after")])
|
||||
assert ch.get() == "after"
|
||||
|
||||
@@ -0,0 +1,478 @@
|
||||
"""Benchmark: DeltaChannel snapshot_frequency — storage vs. read-depth tradeoff.
|
||||
|
||||
Run directly: python tests/test_delta_channel_benchmark.py
|
||||
Run via pytest: pytest tests/test_delta_channel_benchmark.py -s
|
||||
|
||||
Part 1 — baseline (original): DeltaChannel(inf) vs add_messages (BinOp).
|
||||
Part 2 — snapshot_frequency sweep: shows the storage/read-latency tradeoff
|
||||
across frequencies [1, 5, 10, 50, inf] at scale.
|
||||
|
||||
Key insight:
|
||||
snapshot_frequency=inf → O(N) storage, O(N) read depth (pure delta)
|
||||
snapshot_frequency=N → O(N²/N) storage, O(N) read depth bounded by freq
|
||||
snapshot_frequency=1 → O(N²) storage, O(1) read depth (full snapshot)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer, add_messages
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
_POSTGRES_AVAILABLE = True
|
||||
_POSTGRES_URI = os.environ.get(
|
||||
"LANGGRAPH_BENCH_POSTGRES_URI",
|
||||
"postgres://postgres@localhost:5432/postgres?sslmode=disable",
|
||||
)
|
||||
except ImportError:
|
||||
_POSTGRES_AVAILABLE = False
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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. "
|
||||
"We've had prior incidents in this area and want to be deliberate. "
|
||||
"What should we prioritize first, and are there known failure modes we should design around from the start?"
|
||||
)
|
||||
|
||||
_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(_messages_delta_reducer)]
|
||||
|
||||
|
||||
def _make_delta_state(snapshot_frequency: int | float) -> type:
|
||||
"""Create a TypedDict with DeltaChannel at the given snapshot_frequency."""
|
||||
channel = DeltaChannel(
|
||||
_messages_delta_reducer, snapshot_frequency=snapshot_frequency
|
||||
)
|
||||
# Use the functional TypedDict form so the Annotated type is stored as an
|
||||
# already-evaluated object rather than a forward-reference string (which
|
||||
# would fail when get_type_hints tries to resolve 'snapshot_frequency').
|
||||
return TypedDict( # type: ignore[return-value]
|
||||
f"DeltaState_freq{snapshot_frequency}",
|
||||
{"messages": Annotated[list, channel]},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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).
|
||||
Read latency is the average of 5 get_state calls after the full history
|
||||
is built — forces state rehydration including ancestor replay if needed.
|
||||
"""
|
||||
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
|
||||
|
||||
t1 = time.perf_counter()
|
||||
for _ in range(5):
|
||||
graph.get_state(config)
|
||||
read_elapsed = (time.perf_counter() - t1) / 5
|
||||
|
||||
blob_bytes = (
|
||||
_total_blob_bytes(graph.checkpointer)
|
||||
if isinstance(graph.checkpointer, MemorySaver)
|
||||
else -1
|
||||
)
|
||||
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:
|
||||
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"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpointer factories
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _pg_saver(thread_id: str = "bench"):
|
||||
"""Context manager that yields a fresh PostgresSaver and cleans up after."""
|
||||
with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver:
|
||||
saver.setup()
|
||||
with saver._cursor() as cur:
|
||||
for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
|
||||
cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,))
|
||||
yield saver
|
||||
with saver._cursor() as cur:
|
||||
for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
|
||||
cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,))
|
||||
|
||||
|
||||
def _checkpointers() -> list[tuple[str, Any]]:
|
||||
"""Return (label, saver_or_None) pairs for available checkpointers."""
|
||||
result: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
if _POSTGRES_AVAILABLE:
|
||||
try:
|
||||
import psycopg
|
||||
|
||||
psycopg.connect(_POSTGRES_URI).close()
|
||||
result.append(("Postgres", "postgres"))
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part 1: baseline DeltaChannel(inf) vs add_messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
DELTA_ONLY_TURN_COUNTS = [1000]
|
||||
|
||||
|
||||
def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
|
||||
W = 72
|
||||
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
return contextlib.nullcontext(None)
|
||||
return _pg_saver()
|
||||
|
||||
rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = []
|
||||
for turns in BASELINE_TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt))
|
||||
for turns in DELTA_ONLY_TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, None, d_bytes, None, d_rt, None, d_wt))
|
||||
|
||||
def _bytes_or_na(v: Any) -> str:
|
||||
if v is None or v < 0:
|
||||
return "n/a"
|
||||
return _fmt_bytes(v)
|
||||
|
||||
def _ms_or_na(v: Any) -> str:
|
||||
return "n/a" if v is None else f"{v * 1000:.1f}ms"
|
||||
|
||||
print(f"\n [{cp_label}] Storage (blob bytes)")
|
||||
print(
|
||||
f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}"
|
||||
)
|
||||
print(" " + "-" * (W - 2))
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0:
|
||||
ratio_str = "n/a"
|
||||
else:
|
||||
ratio = b_bytes / d_bytes if d_bytes else float("inf")
|
||||
ratio_str = f"{ratio:.0f}x"
|
||||
print(
|
||||
f" {turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} {ratio_str:>8}"
|
||||
)
|
||||
|
||||
print(f"\n [{cp_label}] Read latency (avg of 5 get_state calls)")
|
||||
print(f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12}")
|
||||
print(" " + "-" * (W - 2))
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
print(
|
||||
f" {turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
|
||||
)
|
||||
|
||||
|
||||
def run_baseline_benchmark() -> None:
|
||||
print()
|
||||
print("Part 1 — DeltaChannel(inf) vs add_messages: storage & latency")
|
||||
print("=" * 72)
|
||||
for cp_label, cp_hint in _checkpointers():
|
||||
_run_baseline_for_checkpointer(cp_label, cp_hint)
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Part 2: snapshot_frequency sweep
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Frequencies to test. 1 = always snapshot (like BinOp), inf = pure delta.
|
||||
SNAPSHOT_FREQUENCIES: list[int | float] = [1, 5, 10, 50, math.inf]
|
||||
|
||||
# Turn counts for the sweep — high enough to show storage divergence.
|
||||
SWEEP_TURN_COUNTS = [50, 100, 500]
|
||||
|
||||
|
||||
def _freq_label(freq: int | float) -> str:
|
||||
if freq == math.inf:
|
||||
return "inf"
|
||||
return str(int(freq))
|
||||
|
||||
|
||||
def _run_sweep_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
return contextlib.nullcontext(None)
|
||||
return _pg_saver()
|
||||
|
||||
# Collect results: {turns: {freq_label: (write_s, read_s, bytes)}}
|
||||
results: dict[int, dict[str, tuple[float, float, int]]] = {}
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
results[turns] = {}
|
||||
for freq in SNAPSHOT_FREQUENCIES:
|
||||
state_cls = _make_delta_state(freq)
|
||||
with _make_saver() as saver:
|
||||
wt, rt, bb = _run_turns(turns, state_cls, saver)
|
||||
results[turns][_freq_label(freq)] = (wt, rt, bb)
|
||||
|
||||
freq_labels = [_freq_label(f) for f in SNAPSHOT_FREQUENCIES]
|
||||
col_w = 12
|
||||
|
||||
header = f" {'turns':>6} {'ctx':>10}" + "".join(
|
||||
f" {f'freq={freq_label}':>{col_w}}" for freq_label in freq_labels
|
||||
)
|
||||
|
||||
print(f"\n [{cp_label}] Storage (blob bytes) — lower is better")
|
||||
print(header)
|
||||
print(" " + "-" * (len(header) - 2))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
row = f" {turns:>6} {_approx_tokens(turns):>10}"
|
||||
for label in freq_labels:
|
||||
_, _, bb = results[turns][label]
|
||||
row += f" {_fmt_bytes(bb) if bb >= 0 else 'n/a':>{col_w}}"
|
||||
print(row)
|
||||
|
||||
print(f"\n [{cp_label}] Read latency (avg of 5 get_state) — lower is better")
|
||||
print(header)
|
||||
print(" " + "-" * (len(header) - 2))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
row = f" {turns:>6} {_approx_tokens(turns):>10}"
|
||||
for label in freq_labels:
|
||||
_, rt, _ = results[turns][label]
|
||||
row += f" {f'{rt * 1000:.1f}ms':>{col_w}}"
|
||||
print(row)
|
||||
|
||||
print(
|
||||
f"\n [{cp_label}] Per-invoke write latency (total / turns) — lower is better"
|
||||
)
|
||||
print(header)
|
||||
print(" " + "-" * (len(header) - 2))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
row = f" {turns:>6} {_approx_tokens(turns):>10}"
|
||||
for label in freq_labels:
|
||||
wt, _, _ = results[turns][label]
|
||||
row += f" {f'{(wt / turns) * 1000:.1f}ms':>{col_w}}"
|
||||
print(row)
|
||||
|
||||
|
||||
def run_snapshot_freq_benchmark() -> None:
|
||||
print()
|
||||
print("Part 2 — DeltaChannel snapshot_frequency sweep")
|
||||
print("Lower freq → fewer snapshots → less storage but deeper read replay")
|
||||
print("=" * 80)
|
||||
for cp_label, cp_hint in _checkpointers():
|
||||
_run_sweep_for_checkpointer(cp_label, cp_hint)
|
||||
print()
|
||||
print("Legend:")
|
||||
print(
|
||||
" freq=1 snapshot every write (full blob always — same as add_messages / BinOp)"
|
||||
)
|
||||
print(" freq=N snapshot every N writes; read walks at most N ancestor writes")
|
||||
print(" freq=inf pure delta; read walks entire ancestor chain")
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry points
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
|
||||
)
|
||||
def test_delta_channel_baseline_benchmark(capsys: Any) -> None:
|
||||
"""DeltaChannel(inf) uses less storage than add_messages at scale."""
|
||||
with capsys.disabled():
|
||||
run_baseline_benchmark()
|
||||
|
||||
for turns in [25, 50]:
|
||||
_, _, b_bytes = _run_turns(turns, BinaryState)
|
||||
_, _, d_bytes = _run_turns(turns, DeltaState)
|
||||
assert d_bytes < b_bytes, (
|
||||
f"DeltaChannel should use less storage at {turns} turns, "
|
||||
f"got delta={d_bytes} binary={b_bytes}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
|
||||
)
|
||||
def test_snapshot_freq_benchmark(capsys: Any) -> None:
|
||||
"""snapshot_frequency trades storage for bounded read depth."""
|
||||
with capsys.disabled():
|
||||
run_snapshot_freq_benchmark()
|
||||
|
||||
# Correctness: results at all frequencies should agree on final state.
|
||||
n_turns = 20
|
||||
states: dict[str, list] = {}
|
||||
for freq in SNAPSHOT_FREQUENCIES:
|
||||
state_cls = _make_delta_state(freq)
|
||||
graph = _make_graph(state_cls)
|
||||
config = {"configurable": {"thread_id": "correctness"}}
|
||||
for i in range(n_turns):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
|
||||
config,
|
||||
)
|
||||
state = graph.get_state(config)
|
||||
states[_freq_label(freq)] = [m.id for m in state.values["messages"]]
|
||||
|
||||
ref = states["inf"]
|
||||
for label, msg_ids in states.items():
|
||||
assert msg_ids == ref, (
|
||||
f"freq={label} produced different message IDs than freq=inf"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Script entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_baseline_benchmark()
|
||||
run_snapshot_freq_benchmark()
|
||||
sys.exit(0)
|
||||
@@ -0,0 +1,613 @@
|
||||
"""Tests for the BinaryOperatorAggregate -> DeltaChannel migration path.
|
||||
|
||||
A thread written under `BinaryOperatorAggregate(...)` must keep working
|
||||
after its annotation is swapped to `DeltaChannel(...)` on the same
|
||||
checkpointer — pre-migration state visible at each *settled* ancestor
|
||||
checkpoint is preserved, and post-migration writes fold on top through
|
||||
the reducer.
|
||||
|
||||
Mechanism under test: the saver's `_get_channel_writes_history(config,
|
||||
channel)` walks the parent chain; when it encounters an ancestor whose
|
||||
`channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it
|
||||
returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses
|
||||
it as the base value, and `replay_writes(writes)` folds on-path deltas.
|
||||
|
||||
Scenarios covered:
|
||||
|
||||
1. **Basic migration (sync + async)**: build pre-migration state with
|
||||
`BinaryOperatorAggregate`, swap the annotation to `DeltaChannel` on
|
||||
the same checkpointer, and verify that every settled pre-migration
|
||||
super-step boundary (`next=('__start__',)`) round-trips exactly
|
||||
under the delta-channel view.
|
||||
2. **Time travel into a pre-migration checkpoint** after migration —
|
||||
`graph.get_state(pre_migration_config)` at a settled ancestor
|
||||
returns the same state as under the binop channel.
|
||||
3. **Continuing a migrated thread**: driving one more super-step after
|
||||
migration produces a state that includes the pre-migration settled
|
||||
prefix plus the new delta write — proving `from_checkpoint(seed)` +
|
||||
`replay_writes` correctly fold post-migration deltas onto the
|
||||
pre-migration seed.
|
||||
4. **Base-saver fallback path**: a third-party-style subclass that
|
||||
removes the optimized `InMemorySaver` override and falls back to
|
||||
`BaseCheckpointSaver._get_channel_writes_history` must produce the
|
||||
same result as the optimized path.
|
||||
5. **Channel-type isolation across threads**: two threads on the same
|
||||
checkpointer under the delta-channel graph — one freshly-started,
|
||||
one migrated from pre-migration state — don't cross-contaminate.
|
||||
The parent-chain walk is scoped to the thread.
|
||||
|
||||
TODO: add postgres variants in the existing `libs/checkpoint-postgres`
|
||||
test files (different fixture setup; not this file).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer, add_messages
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factories
|
||||
#
|
||||
# A minimal reducer (`operator.add` on lists of str) with a noop node keeps
|
||||
# state change localized to the HumanMessage-like payload passed through
|
||||
# `invoke`. That isolates the pre/post-migration parity assertions to
|
||||
# channel-hydration semantics.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _noop(_state: Any) -> dict:
|
||||
return {}
|
||||
|
||||
|
||||
def _list_concat(state: list, writes: list) -> list:
|
||||
result = list(state)
|
||||
for w in writes:
|
||||
result.extend(w if isinstance(w, list) else [w])
|
||||
return result
|
||||
|
||||
|
||||
def _binop_graph(checkpointer: Any) -> Any:
|
||||
class BinopState(TypedDict):
|
||||
items: Annotated[list, BinaryOperatorAggregate(list, operator.add)]
|
||||
|
||||
return (
|
||||
StateGraph(BinopState)
|
||||
.add_node("noop", _noop)
|
||||
.add_edge(START, "noop")
|
||||
.add_edge("noop", END)
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
|
||||
|
||||
def _delta_graph(checkpointer: Any) -> Any:
|
||||
class DeltaState(TypedDict):
|
||||
items: Annotated[list, DeltaChannel(_list_concat)]
|
||||
|
||||
return (
|
||||
StateGraph(DeltaState)
|
||||
.add_node("noop", _noop)
|
||||
.add_edge(START, "noop")
|
||||
.add_edge("noop", END)
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
|
||||
|
||||
def _drive(graph: Any, config: dict, tag: str, n: int) -> None:
|
||||
for i in range(n):
|
||||
graph.invoke({"items": [f"{tag}{i}"]}, config)
|
||||
|
||||
|
||||
async def _adrive(graph: Any, config: dict, tag: str, n: int) -> None:
|
||||
for i in range(n):
|
||||
await graph.ainvoke({"items": [f"{tag}{i}"]}, config)
|
||||
|
||||
|
||||
def _settled_boundaries(history: list) -> list[tuple[dict, list]]:
|
||||
"""Return `[(config, items), ...]` for every checkpoint in `history`
|
||||
whose `next == ('__start__',)` — the stable boundaries between invokes.
|
||||
"""
|
||||
return [
|
||||
(s.config, list(s.values.get("items", [])))
|
||||
for s in history
|
||||
if s.next == ("__start__",)
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Basic migration (sync + async)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_basic_migration_preserves_pre_migration_state() -> None:
|
||||
"""Build state under `BinaryOperatorAggregate`, migrate to
|
||||
`DeltaChannel` on the same checkpointer, and verify that every
|
||||
settled pre-migration super-step boundary round-trips exactly.
|
||||
|
||||
Settled boundaries (`next=('__start__',)`) are the stable hydration
|
||||
targets for the migration path: writes that produced the NEXT
|
||||
super-step are kept as `pending_writes` on the ancestor, so walking
|
||||
from a descendant finds the ancestor's blob as the seed and
|
||||
reconstructs the correct state.
|
||||
"""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "basic-sync"}}
|
||||
|
||||
# Pre-migration: accumulate items across 3 invokes.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 3)
|
||||
|
||||
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
|
||||
assert len(pre_boundaries) >= 2, "expected multiple settled boundaries"
|
||||
|
||||
# Migrate: swap the annotation on the same checkpointer.
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
for cfg, items in pre_boundaries:
|
||||
snap = delta.get_state(cfg)
|
||||
assert list(snap.values.get("items", [])) == items, (
|
||||
f"snapshot mismatch at {cfg['configurable']['checkpoint_id']}: "
|
||||
f"expected {items}, got {snap.values.get('items', [])}"
|
||||
)
|
||||
|
||||
|
||||
async def test_basic_migration_preserves_pre_migration_state_async() -> None:
|
||||
"""Async variant of the basic migration scenario."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "basic-async"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
await _adrive(binop, config, "u", 3)
|
||||
|
||||
pre_history = [s async for s in binop.aget_state_history(config)]
|
||||
pre_boundaries = _settled_boundaries(pre_history)
|
||||
assert len(pre_boundaries) >= 2
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
for cfg, items in pre_boundaries:
|
||||
snap = await delta.aget_state(cfg)
|
||||
assert list(snap.values.get("items", [])) == items, (
|
||||
f"async snapshot mismatch at {cfg['configurable']['checkpoint_id']}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Time travel into a pre-migration checkpoint after migration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_time_travel_into_pre_migration_checkpoint() -> None:
|
||||
"""After migration, `graph.get_state(pre_migration_config)` at a
|
||||
settled ancestor returns the state as stored at that point."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "time-travel"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 3)
|
||||
|
||||
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
|
||||
assert pre_boundaries, "no settled ancestors to time-travel to"
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
# Pick the oldest non-empty boundary — a long distance to walk back.
|
||||
non_empty = [(cfg, items) for cfg, items in pre_boundaries if items]
|
||||
assert non_empty, "expected at least one non-empty boundary"
|
||||
target_cfg, expected_items = non_empty[-1]
|
||||
|
||||
snap = delta.get_state(target_cfg)
|
||||
assert list(snap.values.get("items", [])) == expected_items
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Continuing a migrated thread: deltas fold onto pre-migration seed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_continuing_migrated_thread_folds_deltas_on_seed() -> None:
|
||||
"""Resume a pre-migration settled ancestor via `invoke(None, cfg)`
|
||||
under the delta-channel graph. Since the pre-migration checkpoint
|
||||
has an existing `pending_writes` entry (the input for the NEXT
|
||||
super-step), re-running from that ancestor reproduces the same
|
||||
post-ancestor state as the original binop run.
|
||||
|
||||
This proves the seed-terminator + write-replay pipeline works
|
||||
end-to-end across the migration boundary.
|
||||
"""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "continue"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 2)
|
||||
|
||||
# Pick the oldest settled boundary with non-empty state.
|
||||
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
|
||||
target_cfg, seed_items = next(
|
||||
(cfg, items) for cfg, items in reversed(pre_boundaries) if items
|
||||
)
|
||||
assert seed_items, "need a non-empty seed boundary"
|
||||
|
||||
# Migrate and resume from the pre-migration ancestor. `invoke(None,
|
||||
# cfg)` replays the pending writes staged at `cfg` under the new
|
||||
# channel; the reducer folds those deltas onto the seed.
|
||||
delta = _delta_graph(checkpointer)
|
||||
result = delta.invoke(None, target_cfg)
|
||||
|
||||
# The resumed state must include the pre-migration seed items in order.
|
||||
result_items = list(result.get("items", []))
|
||||
for idx, prefix_item in enumerate(seed_items):
|
||||
assert result_items[idx] == prefix_item, (
|
||||
f"pre-migration seed item at {idx} not preserved: "
|
||||
f"got {result_items[: idx + 1]}, expected {seed_items}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Base-saver fallback path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ThirdPartyStyleSaver(InMemorySaver):
|
||||
"""Simulates a third-party saver that inherits the reference
|
||||
`_get_channel_writes_history` implementation from
|
||||
`BaseCheckpointSaver` rather than overriding it.
|
||||
|
||||
We rebind the two methods to the base-class versions (via MRO) so
|
||||
the fallback path is exercised even though the storage layer is
|
||||
still the in-memory one.
|
||||
"""
|
||||
|
||||
# MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...]
|
||||
_get_channel_writes_history = ( # type: ignore[assignment]
|
||||
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
_aget_channel_writes_history = ( # type: ignore[assignment]
|
||||
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
|
||||
def test_base_saver_fallback_matches_optimized_override() -> None:
|
||||
"""The reference `BaseCheckpointSaver` implementation must produce
|
||||
the same migration behavior as the optimized `InMemorySaver`
|
||||
override. We drive the same migration scenario through both savers
|
||||
and assert per-snapshot parity in the delta-channel view."""
|
||||
|
||||
# Fast path: optimized InMemorySaver override.
|
||||
fast_saver = InMemorySaver()
|
||||
fast_config = {"configurable": {"thread_id": "fast"}}
|
||||
fast_binop = _binop_graph(fast_saver)
|
||||
_drive(fast_binop, fast_config, "u", 3)
|
||||
fast_delta = _delta_graph(fast_saver)
|
||||
fast_history = [
|
||||
(s.next, list(s.values.get("items", [])))
|
||||
for s in fast_delta.get_state_history(fast_config)
|
||||
]
|
||||
|
||||
# Slow path: base-class fallback.
|
||||
slow_saver = _ThirdPartyStyleSaver()
|
||||
slow_config = {"configurable": {"thread_id": "slow"}}
|
||||
slow_binop = _binop_graph(slow_saver)
|
||||
_drive(slow_binop, slow_config, "u", 3)
|
||||
slow_delta = _delta_graph(slow_saver)
|
||||
slow_history = [
|
||||
(s.next, list(s.values.get("items", [])))
|
||||
for s in slow_delta.get_state_history(slow_config)
|
||||
]
|
||||
|
||||
assert slow_history == fast_history, (
|
||||
"base-saver fallback should match optimized-override behavior; "
|
||||
f"fast={fast_history}, slow={slow_history}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Thread isolation under mixed-generation storage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delta_and_migrated_threads_do_not_cross_contaminate() -> None:
|
||||
"""Two threads sharing a checkpointer — one migrated from
|
||||
pre-migration state, one freshly-started under DeltaChannel — must
|
||||
maintain independent state. The parent-chain walk in
|
||||
`_get_channel_writes_history` must be scoped to the target thread.
|
||||
"""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
migrated_cfg = {"configurable": {"thread_id": "migrated"}}
|
||||
fresh_cfg = {"configurable": {"thread_id": "fresh"}}
|
||||
|
||||
# Thread A: pre-migration build-up.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, migrated_cfg, "m", 2)
|
||||
|
||||
# Thread B: fresh delta-channel run.
|
||||
delta = _delta_graph(checkpointer)
|
||||
_drive(delta, fresh_cfg, "f", 2)
|
||||
|
||||
# Thread A: migrate and confirm its state is anchored in its own
|
||||
# thread's pre-migration history (tag 'm'), never mixing in tag 'f'.
|
||||
migrated_boundaries = _settled_boundaries(
|
||||
list(delta.get_state_history(migrated_cfg))
|
||||
)
|
||||
assert migrated_boundaries, "migrated thread has no settled boundaries"
|
||||
for _, items in migrated_boundaries:
|
||||
for it in items:
|
||||
assert it.startswith("m"), (
|
||||
f"migrated thread leaked item from other thread: {it}"
|
||||
)
|
||||
|
||||
# Thread B: settled boundaries must only contain 'f' tags.
|
||||
fresh_boundaries = _settled_boundaries(list(delta.get_state_history(fresh_cfg)))
|
||||
assert fresh_boundaries, "fresh thread has no settled boundaries"
|
||||
for _, items in fresh_boundaries:
|
||||
for it in items:
|
||||
assert it.startswith("f"), (
|
||||
f"fresh thread leaked item from migrated thread: {it}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. Tip-of-pre-migration hydration: the latest checkpoint from a binop-run
|
||||
# thread has a real accumulated value in its own `channel_values["items"]`.
|
||||
# When hydrated under the delta-channel graph via `get_state(config)` with no
|
||||
# `checkpoint_id`, the short-circuit must use that value directly instead of
|
||||
# walking ancestors (which would skip the tip's own blob).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_tip_of_pre_migration_hydrates_directly() -> None:
|
||||
"""`graph.get_state(config)` at the latest (pre-migration) checkpoint
|
||||
returns the full accumulated list stored in that checkpoint's own
|
||||
`channel_values`. The hydration must not walk ancestors past it."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "tip-sync"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 3)
|
||||
|
||||
binop_tip = binop.get_state(config)
|
||||
expected_items = list(binop_tip.values.get("items", []))
|
||||
assert expected_items == ["u0", "u1", "u2"], (
|
||||
f"sanity: pre-migration tip should accumulate all 3 items, got {expected_items}"
|
||||
)
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
snap = delta.get_state(config)
|
||||
assert list(snap.values.get("items", [])) == expected_items, (
|
||||
f"tip hydration mismatch: expected {expected_items}, "
|
||||
f"got {snap.values.get('items', [])}"
|
||||
)
|
||||
|
||||
|
||||
async def test_tip_of_pre_migration_hydrates_directly_async() -> None:
|
||||
"""Async variant of the tip-of-pre-migration hydration scenario."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "tip-async"}}
|
||||
|
||||
binop = _binop_graph(checkpointer)
|
||||
await _adrive(binop, config, "u", 3)
|
||||
|
||||
binop_tip = await binop.aget_state(config)
|
||||
expected_items = list(binop_tip.values.get("items", []))
|
||||
assert expected_items == ["u0", "u1", "u2"]
|
||||
|
||||
delta = _delta_graph(checkpointer)
|
||||
|
||||
snap = await delta.aget_state(config)
|
||||
assert list(snap.values.get("items", [])) == expected_items, (
|
||||
f"async tip hydration mismatch: expected {expected_items}, "
|
||||
f"got {snap.values.get('items', [])}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. `update_state` after migration writes a real value to the new
|
||||
# checkpoint's `channel_values` (not a sentinel). Hydration must use it
|
||||
# directly — the ancestor walk would skip this blob and return stale state.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_update_state_after_migration_uses_written_value() -> None:
|
||||
"""After migrating and running at least one post-migration super-step
|
||||
(so the thread's tip has a `DELTA_SENTINEL`), `update_state` writes a
|
||||
concrete value to a new checkpoint's `channel_values`. `get_state`
|
||||
must reflect that concrete value."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "update-state"}}
|
||||
|
||||
# Pre-migration: accumulate a little state.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 2)
|
||||
|
||||
# Migrate and run one more super-step so the tip is a post-migration
|
||||
# checkpoint with `DELTA_SENTINEL` in its own `channel_values`.
|
||||
delta = _delta_graph(checkpointer)
|
||||
delta.invoke({"items": ["post"]}, config)
|
||||
|
||||
# `update_state` writes a concrete value into a new checkpoint's blob
|
||||
# via the reducer against the hydrated prior state.
|
||||
delta.update_state(config, {"items": ["x", "y"]})
|
||||
|
||||
snap = delta.get_state(config)
|
||||
updated_items = list(snap.values.get("items", []))
|
||||
# Must include the "x","y" update; without the hydration fix, the
|
||||
# update_state-written blob would be skipped in favor of an ancestor
|
||||
# walk, and the update values would disappear.
|
||||
assert "x" in updated_items and "y" in updated_items, (
|
||||
f"update_state values missing from snapshot: {updated_items}"
|
||||
)
|
||||
# The "x","y" items should be folded onto the prior accumulated state,
|
||||
# not stand alone. This verifies the update-written blob is used
|
||||
# directly by `get_state` (no ancestor walk past it).
|
||||
assert len(updated_items) >= 4, (
|
||||
f"update_state snapshot should preserve pre-update state, got {updated_items}"
|
||||
)
|
||||
assert updated_items[-2:] == ["x", "y"], (
|
||||
f"update_state deltas should be at the tail, got {updated_items}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 8. Fork from an `update_state` checkpoint: a new run branched off the
|
||||
# update_state-produced checkpoint must see that checkpoint's concrete
|
||||
# `channel_values` as its base, with new deltas folded on top.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fork_from_update_state_checkpoint() -> None:
|
||||
"""Branching a new run from the checkpoint produced by `update_state`
|
||||
must use that checkpoint's concrete blob as the base. Additional
|
||||
deltas from the forked run fold onto it through the reducer."""
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "fork"}}
|
||||
|
||||
# Pre-migration build-up, then migrate and add one post-migration step.
|
||||
binop = _binop_graph(checkpointer)
|
||||
_drive(binop, config, "u", 2)
|
||||
delta = _delta_graph(checkpointer)
|
||||
delta.invoke({"items": ["post"]}, config)
|
||||
|
||||
# Apply `update_state` and capture the returned config (references
|
||||
# the new checkpoint produced by the update).
|
||||
update_cfg = delta.update_state(config, {"items": ["x", "y"]})
|
||||
|
||||
update_snap = delta.get_state(update_cfg)
|
||||
base_items = list(update_snap.values.get("items", []))
|
||||
assert "x" in base_items and "y" in base_items, (
|
||||
f"update_state values missing from snapshot: {base_items}"
|
||||
)
|
||||
assert base_items[-2:] == ["x", "y"], (
|
||||
f"sanity: update_state deltas should be at the tail, got {base_items}"
|
||||
)
|
||||
|
||||
# Fork: invoke from the update_state checkpoint with a new delta.
|
||||
forked = delta.invoke({"items": ["fork0"]}, update_cfg)
|
||||
forked_items = list(forked.get("items", []))
|
||||
# The fork must see the update_state-written blob as its base (not
|
||||
# walk past it), and the new delta must fold on top of it.
|
||||
assert forked_items[: len(base_items)] == base_items, (
|
||||
f"fork lost update_state base: base={base_items}, forked={forked_items}"
|
||||
)
|
||||
assert forked_items[-1] == "fork0", f"fork delta not appended: {forked_items}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 9. Migration from `add_messages` → `DeltaChannel(_messages_delta_reducer)`
|
||||
#
|
||||
# `add_messages` is the primary real-world use case: it creates a
|
||||
# BinaryOperatorAggregate with dedup-by-ID and RemoveMessage semantics.
|
||||
# After swapping the annotation to DeltaChannel, pre-migration blobs
|
||||
# (plain lists of Message objects) must be used directly as the seed.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _add_messages_graph(checkpointer: Any) -> Any:
|
||||
class MessagesState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("noop", _noop)
|
||||
.add_edge(START, "noop")
|
||||
.add_edge("noop", END)
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
|
||||
|
||||
def _delta_messages_graph(checkpointer: Any) -> Any:
|
||||
class DeltaMessagesState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
return (
|
||||
StateGraph(DeltaMessagesState)
|
||||
.add_node("noop", _noop)
|
||||
.add_edge(START, "noop")
|
||||
.add_edge("noop", END)
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
|
||||
|
||||
def test_add_messages_to_delta_migration_preserves_message_history() -> None:
|
||||
"""Migration from `add_messages` to `DeltaChannel(_messages_delta_reducer)`
|
||||
preserves message ordering and IDs at both the tip and settled ancestor
|
||||
boundaries.
|
||||
|
||||
The pre-migration blob is a plain list of Message objects; DeltaChannel
|
||||
must use it directly as the seed without walking ancestors past it.
|
||||
"""
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "add-messages-migration"}}
|
||||
|
||||
pre_graph = _add_messages_graph(checkpointer)
|
||||
pre_graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
||||
pre_graph.invoke({"messages": [AIMessage(content="hi", id="a1")]}, config)
|
||||
pre_graph.invoke({"messages": [HumanMessage(content="thanks", id="h2")]}, config)
|
||||
|
||||
pre_tip = pre_graph.get_state(config)
|
||||
assert [m.id for m in pre_tip.values["messages"]] == ["h1", "a1", "h2"]
|
||||
|
||||
delta_graph = _delta_messages_graph(checkpointer)
|
||||
|
||||
# Tip: latest checkpoint has a full list blob — must use it directly.
|
||||
snap = delta_graph.get_state(config)
|
||||
assert [m.id for m in snap.values["messages"]] == ["h1", "a1", "h2"], (
|
||||
f"tip hydration mismatch: got {[m.id for m in snap.values['messages']]}"
|
||||
)
|
||||
|
||||
# Settled ancestor boundaries must also match.
|
||||
pre_settled = [
|
||||
[m.id for m in s.values.get("messages", [])]
|
||||
for s in pre_graph.get_state_history(config)
|
||||
if s.next == ("__start__",)
|
||||
]
|
||||
delta_settled = [
|
||||
[m.id for m in s.values.get("messages", [])]
|
||||
for s in delta_graph.get_state_history(config)
|
||||
if s.next == ("__start__",)
|
||||
]
|
||||
assert delta_settled == pre_settled, (
|
||||
f"settled boundary mismatch after migration: "
|
||||
f"pre={pre_settled}, delta={delta_settled}"
|
||||
)
|
||||
|
||||
|
||||
async def test_add_messages_to_delta_migration_preserves_message_history_async() -> (
|
||||
None
|
||||
):
|
||||
"""Async variant of the add_messages migration test."""
|
||||
checkpointer = InMemorySaver()
|
||||
config = {"configurable": {"thread_id": "add-messages-migration-async"}}
|
||||
|
||||
pre_graph = _add_messages_graph(checkpointer)
|
||||
await pre_graph.ainvoke(
|
||||
{"messages": [HumanMessage(content="hello", id="h1")]}, config
|
||||
)
|
||||
await pre_graph.ainvoke({"messages": [AIMessage(content="hi", id="a1")]}, config)
|
||||
|
||||
delta_graph = _delta_messages_graph(checkpointer)
|
||||
snap = await delta_graph.aget_state(config)
|
||||
assert [m.id for m in snap.values["messages"]] == ["h1", "a1"], (
|
||||
f"async tip hydration mismatch: got {[m.id for m in snap.values['messages']]}"
|
||||
)
|
||||
@@ -16,7 +16,7 @@ from typing import Annotated, Any, Literal, get_type_hints
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models import GenericFakeChatModel
|
||||
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
|
||||
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, RemoveMessage
|
||||
from langchain_core.runnables import (
|
||||
RunnableConfig,
|
||||
RunnableLambda,
|
||||
@@ -25,6 +25,7 @@ from langchain_core.runnables import (
|
||||
from langchain_core.runnables.graph import Edge
|
||||
from langgraph.cache.base import BaseCache
|
||||
from langgraph.checkpoint.base import (
|
||||
DELTA_SENTINEL,
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
@@ -41,6 +42,7 @@ from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.topic import Topic
|
||||
@@ -49,7 +51,7 @@ from langgraph.config import get_stream_writer
|
||||
from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import MessagesState, add_messages
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
|
||||
from langgraph.pregel import (
|
||||
NodeBuilder,
|
||||
Pregel,
|
||||
@@ -9400,3 +9402,254 @@ 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."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
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."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
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."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
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."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
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"
|
||||
|
||||
|
||||
async def test_delta_channel_durability_exit_stores_snapshot() -> None:
|
||||
"""DeltaChannel must reload from a durability='exit' checkpoint."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
return {"messages": [AIMessage(content="reply", id="ai1")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "delta-exit-test"}}
|
||||
|
||||
result = graph.invoke(
|
||||
{"messages": [HumanMessage(content="hello", id="h1")]},
|
||||
config,
|
||||
durability="exit",
|
||||
)
|
||||
assert [m.content for m in result["messages"]] == ["hello", "reply"]
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert [m.content for m in state.values["messages"]] == ["hello", "reply"]
|
||||
|
||||
|
||||
async def test_delta_channel_async_write_ordering() -> None:
|
||||
"""In async mode, DeltaChannel write futures are awaited before the checkpoint
|
||||
is committed, so aput_writes always precedes aput for sentinel checkpoints."""
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
i = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"r{i}", id=f"ai{i}")]}
|
||||
|
||||
order: list[str] = []
|
||||
original_aput_writes = InMemorySaver.aput_writes
|
||||
original_aput = InMemorySaver.aput
|
||||
|
||||
async def tracked_aput_writes(self, config, writes, task_id, task_path=""):
|
||||
result = await original_aput_writes(self, config, writes, task_id, task_path)
|
||||
order.append("aput_writes")
|
||||
return result
|
||||
|
||||
async def tracked_aput(self, config, checkpoint, metadata, new_versions):
|
||||
has_sentinel = any(
|
||||
v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values()
|
||||
)
|
||||
order.append("aput_sentinel" if has_sentinel else "aput_other")
|
||||
return await original_aput(self, config, checkpoint, metadata, new_versions)
|
||||
|
||||
InMemorySaver.aput_writes = tracked_aput_writes
|
||||
InMemorySaver.aput = tracked_aput
|
||||
try:
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "async-ordering-test"}}
|
||||
|
||||
for i in range(3):
|
||||
await graph.ainvoke(
|
||||
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config
|
||||
)
|
||||
|
||||
# Every aput_sentinel must be preceded by at least one aput_writes
|
||||
for i, event in enumerate(order):
|
||||
if event == "aput_sentinel":
|
||||
preceding = order[:i]
|
||||
assert "aput_writes" in preceding, (
|
||||
f"aput_sentinel at {i} had no preceding aput_writes: {order}"
|
||||
)
|
||||
last_write_idx = max(
|
||||
j for j, e in enumerate(order[:i]) if e == "aput_writes"
|
||||
)
|
||||
assert last_write_idx < i, (
|
||||
f"aput_writes at {last_write_idx} should precede aput_sentinel at {i}: {order}"
|
||||
)
|
||||
finally:
|
||||
InMemorySaver.aput_writes = original_aput_writes
|
||||
InMemorySaver.aput = original_aput
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
assert len(state.values["messages"]) == 6 # 3 human + 3 AI
|
||||
|
||||
@@ -6101,6 +6101,36 @@ async def test_parent_command(
|
||||
)
|
||||
|
||||
|
||||
async def test_delta_channel_durability_exit_stores_snapshot_async() -> None:
|
||||
"""DeltaChannel must reload from an async durability='exit' checkpoint."""
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
async def respond(state: State) -> dict:
|
||||
return {"messages": [AIMessage(content="reply", id="ai1")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "delta-exit-async-test"}}
|
||||
|
||||
result = await graph.ainvoke(
|
||||
{"messages": [HumanMessage(content="hello", id="h1")]},
|
||||
config,
|
||||
durability="exit",
|
||||
)
|
||||
assert [m.content for m in result["messages"]] == ["hello", "reply"]
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
assert [m.content for m in state.values["messages"]] == ["hello", "reply"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_interrupt_subgraph(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
class State(TypedDict):
|
||||
|
||||
Generated
+30
-14
@@ -7,6 +7,9 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
|
||||
[options]
|
||||
prerelease-mode = "allow"
|
||||
|
||||
[[package]]
|
||||
name = "aiosqlite"
|
||||
version = "0.22.1"
|
||||
@@ -1348,10 +1351,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
{ name = "langchain-protocol" },
|
||||
{ name = "langsmith" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pydantic" },
|
||||
@@ -1360,14 +1364,26 @@ 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/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
|
||||
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/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langchain-protocol"
|
||||
version = "0.0.14"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/05/bf/efb5e2ed832e4d6d45590e25a9e5191986b291b543bc6a807b48bee070b0/langchain_protocol-0.0.14.tar.gz", hash = "sha256:bc1e8553122e6ede310280462d5813023a172ff2785ccbbdec54d43f3a15e5f2", size = 5862, upload-time = "2026-04-29T16:40:18.657Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c2/e9/06c47ecb2aff08f83dfa30058da3bf86be64862c19569043ed5331bbeecd/langchain_protocol-0.0.14-py3-none-any.whl", hash = "sha256:ffc35089779bd8ca217015180cef5e660fc3b074efdaa0f2e95df73583f1a047", size = 6984, upload-time = "2026-04-29T16:40:17.841Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.2.0a2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1439,7 +1455,7 @@ test = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
@@ -1548,7 +1564,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.0.3"
|
||||
version = "4.1.0a2"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1742,7 +1758,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.13"
|
||||
version = "1.0.10"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1751,7 +1767,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 +2156,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "nbconvert"
|
||||
version = "7.17.1"
|
||||
version = "7.17.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "beautifulsoup4" },
|
||||
@@ -2158,9 +2174,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 +3034,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]]
|
||||
|
||||
@@ -44,7 +44,7 @@ import inspect
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from copy import copy, deepcopy
|
||||
from dataclasses import dataclass, field, replace
|
||||
from dataclasses import dataclass, replace
|
||||
from types import UnionType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
@@ -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"],
|
||||
@@ -836,7 +835,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"],
|
||||
@@ -860,30 +859,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 +876,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 +906,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 +915,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 +947,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 +988,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 +1059,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 +1068,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 +1100,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 +1141,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 +1273,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 +1404,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 +1458,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), ...]}, ...)`'
|
||||
@@ -1722,9 +1621,9 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
|
||||
context: ContextT
|
||||
config: RunnableConfig
|
||||
stream_writer: StreamWriter
|
||||
tools: list[BaseTool]
|
||||
tool_call_id: str | None
|
||||
store: BaseStore | None
|
||||
tools: list[BaseTool] = field(default_factory=list)
|
||||
execution_info: ExecutionInfo | None = None
|
||||
server_info: ServerInfo | None = None
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.0.13"
|
||||
version = "1.0.10"
|
||||
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]
|
||||
|
||||
@@ -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,19 +2016,6 @@ 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_defaults_tools_to_empty_list() -> None:
|
||||
runtime = ToolRuntime(
|
||||
state={},
|
||||
context=None,
|
||||
config={},
|
||||
stream_writer=lambda *args, **kwargs: None,
|
||||
tool_call_id=None,
|
||||
store=None,
|
||||
)
|
||||
|
||||
assert runtime.tools == []
|
||||
|
||||
|
||||
def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
|
||||
"""Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
@@ -2236,195 +2223,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
+7
-7
@@ -249,7 +249,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -261,14 +261,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/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/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/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/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.1.9"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -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.13"
|
||||
version = "1.0.10"
|
||||
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
+7
-7
@@ -262,7 +262,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
@@ -274,14 +274,14 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/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/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/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/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.10"
|
||||
version = "1.1.9"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -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.13"
|
||||
version = "1.0.10"
|
||||
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