mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-28 18:59:42 +02:00
Compare commits
50
Commits
@@ -100,3 +100,4 @@ dmypy.json
|
||||
.turbo
|
||||
.editorconfig
|
||||
.scratch
|
||||
.worktrees/
|
||||
|
||||
@@ -4,7 +4,7 @@ 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 (
|
||||
@@ -32,7 +32,7 @@ Conn = _internal.Conn # For backward compatibility
|
||||
class PostgresSaver(BasePostgresSaver):
|
||||
"""Checkpointer that stores checkpoints in a Postgres database."""
|
||||
|
||||
lock: threading.Lock
|
||||
lock: threading.RLock
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -48,7 +48,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
|
||||
self.conn = conn
|
||||
self.pipe = pipe
|
||||
self.lock = threading.Lock()
|
||||
self.lock = threading.RLock()
|
||||
self.supports_pipeline = Capabilities().has_pipeline()
|
||||
|
||||
@classmethod
|
||||
@@ -442,6 +442,22 @@ class PostgresSaver(BasePostgresSaver):
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
channel_values = self._load_blobs(value["channel_values"])
|
||||
if any(isinstance(v, DeltaChannelSentinel) for v in channel_values.values()):
|
||||
cp_config = cast(
|
||||
RunnableConfig,
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
)
|
||||
with self._cursor() as cur:
|
||||
self._resolve_delta_channels(cp_config, channel_values, cur)
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
@@ -454,7 +470,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
**channel_values,
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
|
||||
@@ -13,6 +13,7 @@ from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
DeltaChannelSentinel,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -391,6 +392,58 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _aget_channel_writes_cur(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
cur: Any,
|
||||
) -> list[Any]:
|
||||
"""Async version of _get_channel_writes_cur — see sync version for rationale."""
|
||||
await cur.execute(
|
||||
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
|
||||
"WHERE thread_id = %s AND checkpoint_ns = %s",
|
||||
(thread_id, checkpoint_ns),
|
||||
)
|
||||
parent_map: dict[str, str | None] = {
|
||||
row["checkpoint_id"]: row["parent_checkpoint_id"]
|
||||
for row in await cur.fetchall()
|
||||
}
|
||||
ancestor_ids: list[str] = []
|
||||
cid: str | None = parent_map.get(checkpoint_id)
|
||||
while cid is not None:
|
||||
ancestor_ids.append(cid)
|
||||
cid = parent_map.get(cid)
|
||||
if not ancestor_ids:
|
||||
return []
|
||||
await cur.execute(
|
||||
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
|
||||
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
||||
" AND checkpoint_id = ANY(%s) "
|
||||
"ORDER BY task_id, idx",
|
||||
(thread_id, checkpoint_ns, channel, ancestor_ids),
|
||||
)
|
||||
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
|
||||
for row in await cur.fetchall():
|
||||
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
|
||||
result = []
|
||||
for cid in reversed(ancestor_ids):
|
||||
for type_tag, blob in writes_by_cp.get(cid, []):
|
||||
result.append(self.serde.loads_typed((type_tag, blob)))
|
||||
return result
|
||||
|
||||
async def aget_channel_writes(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> list[Any]:
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = config["configurable"]["checkpoint_id"]
|
||||
async with self._cursor() as cur:
|
||||
return await self._aget_channel_writes_cur(
|
||||
thread_id, checkpoint_ns, checkpoint_id, channel, cur
|
||||
)
|
||||
|
||||
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
@@ -403,11 +456,31 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
thread_id = value["thread_id"]
|
||||
checkpoint_ns = value["checkpoint_ns"]
|
||||
checkpoint_id = value["checkpoint_id"]
|
||||
blob_values = value["channel_values"]
|
||||
|
||||
channel_values: dict[str, Any] = {}
|
||||
if blob_values:
|
||||
channel_values = self._load_blobs(blob_values)
|
||||
delta_channels = [
|
||||
ch
|
||||
for ch, v in channel_values.items()
|
||||
if isinstance(v, DeltaChannelSentinel)
|
||||
]
|
||||
if delta_channels:
|
||||
async with self._cursor() as cur:
|
||||
for channel in delta_channels:
|
||||
channel_values[channel] = await self._aget_channel_writes_cur(
|
||||
thread_id, checkpoint_ns, checkpoint_id, channel, cur
|
||||
)
|
||||
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
@@ -415,15 +488,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
**channel_values,
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": value["parent_checkpoint_id"],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import random
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from collections.abc import Sequence
|
||||
from importlib.metadata import version as get_version
|
||||
from typing import Any, cast
|
||||
@@ -11,6 +12,7 @@ from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
DeltaChannelSentinel,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
@@ -185,15 +187,78 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
)
|
||||
|
||||
def _load_blobs(
|
||||
self, blob_values: list[tuple[bytes, bytes, bytes]]
|
||||
self,
|
||||
blob_values: Any,
|
||||
) -> dict[str, Any]:
|
||||
if not blob_values:
|
||||
return {}
|
||||
return {
|
||||
k.decode(): self.serde.loads_typed((t.decode(), v))
|
||||
for k, t, v in blob_values
|
||||
if t.decode() != "empty"
|
||||
result: dict[str, Any] = {}
|
||||
for k, t, v in blob_values:
|
||||
type_tag = t.decode()
|
||||
if type_tag != "empty":
|
||||
result[k.decode()] = self.serde.loads_typed((type_tag, v))
|
||||
return result
|
||||
|
||||
def _resolve_delta_channels(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
channel_values: dict[str, Any],
|
||||
cur: Any,
|
||||
) -> None:
|
||||
for channel, value in list(channel_values.items()):
|
||||
if isinstance(value, DeltaChannelSentinel):
|
||||
channel_values[channel] = self._get_channel_writes_cur(
|
||||
config["configurable"]["thread_id"],
|
||||
config["configurable"].get("checkpoint_ns", ""),
|
||||
config["configurable"]["checkpoint_id"],
|
||||
channel,
|
||||
cur,
|
||||
)
|
||||
|
||||
def _get_channel_writes_cur(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
cur: Any,
|
||||
) -> list[Any]:
|
||||
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest.
|
||||
|
||||
Two queries:
|
||||
1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread — cheap, IDs only.
|
||||
2. Walk the ancestor chain in Python, then fetch writes with a plain ANY() filter.
|
||||
"""
|
||||
cur.execute(
|
||||
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
|
||||
"WHERE thread_id = %s AND checkpoint_ns = %s",
|
||||
(thread_id, checkpoint_ns),
|
||||
)
|
||||
parent_map: dict[str, str | None] = {
|
||||
row["checkpoint_id"]: row["parent_checkpoint_id"] for row in cur.fetchall()
|
||||
}
|
||||
ancestor_ids: list[str] = []
|
||||
cid: str | None = parent_map.get(checkpoint_id)
|
||||
while cid is not None:
|
||||
ancestor_ids.append(cid)
|
||||
cid = parent_map.get(cid)
|
||||
if not ancestor_ids:
|
||||
return []
|
||||
cur.execute(
|
||||
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
|
||||
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
||||
" AND checkpoint_id = ANY(%s) "
|
||||
"ORDER BY task_id, idx",
|
||||
(thread_id, checkpoint_ns, channel, ancestor_ids),
|
||||
)
|
||||
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
|
||||
for row in cur.fetchall():
|
||||
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
|
||||
result = []
|
||||
for cid in reversed(ancestor_ids):
|
||||
for type_tag, blob in writes_by_cp.get(cid, []):
|
||||
result.append(self.serde.loads_typed((type_tag, blob)))
|
||||
return result
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
|
||||
@@ -371,3 +371,47 @@ async def test_get_checkpoint_no_channel_values(
|
||||
|
||||
checkpoint = await saver.aget_tuple(config)
|
||||
assert checkpoint.checkpoint["channel_values"] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
|
||||
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
|
||||
pytest.importorskip(
|
||||
"langgraph.channels.delta", reason="langgraph core not installed"
|
||||
)
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
|
||||
async with _saver(saver_name) as saver:
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
config = {"configurable": {"thread_id": "diff-channel-test-1"}}
|
||||
|
||||
await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
|
||||
await graph.ainvoke(
|
||||
{"messages": [HumanMessage(content="there", id="h2")]}, config
|
||||
)
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
msgs = state.values["messages"]
|
||||
assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].content == "hi"
|
||||
assert msgs[1].content == "reply-1"
|
||||
assert msgs[2].content == "there"
|
||||
assert msgs[3].content == "reply-3"
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import dataclasses
|
||||
import logging
|
||||
import threading
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
|
||||
from typing import ( # noqa: UP035
|
||||
Any,
|
||||
Generic,
|
||||
List,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
TypedDict,
|
||||
@@ -28,6 +31,21 @@ from langgraph.checkpoint.serde.types import (
|
||||
|
||||
V = TypeVar("V", int, float, str)
|
||||
PendingWrite = tuple[str, str, Any]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DeltaChannelSentinel:
|
||||
"""Marker stored in checkpoint_blobs for a DeltaChannel field.
|
||||
|
||||
No data is stored here — the actual per-step writes live in checkpoint_writes
|
||||
and are replayed through the reducer at load time.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
_DELTA_RECONSTRUCTION: threading.local = threading.local()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -457,6 +475,56 @@ class BaseCheckpointSaver(Generic[V]):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006
|
||||
"""Collect all writes for `channel` across this checkpoint's ancestry, oldest→newest.
|
||||
|
||||
Default implementation walks the full thread history via `list()`. Savers can
|
||||
override with a more efficient query (InMemorySaver and PostgresSaver do this).
|
||||
"""
|
||||
# Guard against re-entrant calls: when list() triggers reconstruction which
|
||||
# calls list() again, the inner call returns tuples with DeltaChannelSentinel
|
||||
# in channel_values (which get_channel_writes ignores — it only reads
|
||||
# pending_writes). This breaks the recursion safely.
|
||||
if getattr(_DELTA_RECONSTRUCTION, "active", False):
|
||||
return []
|
||||
_DELTA_RECONSTRUCTION.active = True
|
||||
try:
|
||||
result: list[Any] = []
|
||||
target_id = config["configurable"].get("checkpoint_id")
|
||||
for tup in self.list(config):
|
||||
if tup.config["configurable"].get("checkpoint_id") == target_id:
|
||||
continue # skip the checkpoint itself; we want its ancestors' writes
|
||||
if tup.pending_writes:
|
||||
for _, ch, value in tup.pending_writes:
|
||||
if ch == channel:
|
||||
result.append(value)
|
||||
result.reverse() # list() yields newest→oldest; we want oldest→newest
|
||||
return result
|
||||
finally:
|
||||
_DELTA_RECONSTRUCTION.active = False
|
||||
|
||||
async def aget_channel_writes(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> List[Any]: # noqa: UP006
|
||||
"""Async version of get_channel_writes."""
|
||||
if getattr(_DELTA_RECONSTRUCTION, "active", False):
|
||||
return []
|
||||
_DELTA_RECONSTRUCTION.active = True
|
||||
try:
|
||||
result: list[Any] = []
|
||||
target_id = config["configurable"].get("checkpoint_id")
|
||||
async for tup in self.alist(config):
|
||||
if tup.config["configurable"].get("checkpoint_id") == target_id:
|
||||
continue
|
||||
if tup.pending_writes:
|
||||
for _, ch, value in tup.pending_writes:
|
||||
if ch == channel:
|
||||
result.append(value)
|
||||
result.reverse()
|
||||
return result
|
||||
finally:
|
||||
_DELTA_RECONSTRUCTION.active = False
|
||||
|
||||
def get_next_version(self, current: V | None, channel: None) -> V:
|
||||
"""Generate the next version ID for a channel.
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ from collections import defaultdict
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
|
||||
from types import TracebackType
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
@@ -20,6 +20,7 @@ from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
DeltaChannelSentinel,
|
||||
SerializerProtocol,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
@@ -121,16 +122,60 @@ 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 _resolve_delta_channels(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
channel_values: dict[str, Any],
|
||||
) -> None:
|
||||
"""Replace DeltaChannelSentinel entries with reconstructed write lists."""
|
||||
for channel, value in list(channel_values.items()):
|
||||
if isinstance(value, DeltaChannelSentinel):
|
||||
channel_values[channel] = self.get_channel_writes(config, channel)
|
||||
|
||||
def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]:
|
||||
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 collecting checkpoint IDs.
|
||||
chain: list[str] = []
|
||||
current: str | None = checkpoint_id
|
||||
while current is not None:
|
||||
entry = ns_storage.get(current)
|
||||
if entry is None:
|
||||
break
|
||||
chain.append(current)
|
||||
_, _, parent = entry
|
||||
current = parent
|
||||
# Collect writes oldest→newest.
|
||||
result: list[Any] = []
|
||||
for cp_id in reversed(chain):
|
||||
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
|
||||
for (_task_id, _idx), (_, ch, serialized, _) in sorted(step_writes.items()):
|
||||
if ch == channel:
|
||||
result.append(self.serde.loads_typed(serialized))
|
||||
return result
|
||||
|
||||
async def aget_channel_writes(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> list[Any]:
|
||||
return self.get_channel_writes(config, channel)
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
"""Get a checkpoint tuple from the in-memory storage.
|
||||
@@ -153,13 +198,17 @@ class InMemorySaver(
|
||||
checkpoint, metadata, parent_checkpoint_id = saved
|
||||
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
|
||||
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
|
||||
channel_values = self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
)
|
||||
self._resolve_delta_channels(config, channel_values)
|
||||
return CheckpointTuple(
|
||||
config=config,
|
||||
checkpoint={
|
||||
**checkpoint_,
|
||||
"channel_values": self._load_blobs(
|
||||
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
|
||||
),
|
||||
"channel_values": channel_values,
|
||||
},
|
||||
metadata=self.serde.loads_typed(metadata),
|
||||
pending_writes=[
|
||||
@@ -183,19 +232,27 @@ class InMemorySaver(
|
||||
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
|
||||
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
|
||||
checkpoint_ = self.serde.loads_typed(checkpoint)
|
||||
return CheckpointTuple(
|
||||
config={
|
||||
resolved_config = cast(
|
||||
RunnableConfig,
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
)
|
||||
channel_values = self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
)
|
||||
self._resolve_delta_channels(resolved_config, channel_values)
|
||||
return CheckpointTuple(
|
||||
config=resolved_config,
|
||||
checkpoint={
|
||||
**checkpoint_,
|
||||
"channel_values": self._load_blobs(
|
||||
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
|
||||
),
|
||||
"channel_values": channel_values,
|
||||
},
|
||||
metadata=self.serde.loads_typed(metadata),
|
||||
pending_writes=[
|
||||
@@ -290,21 +347,28 @@ class InMemorySaver(
|
||||
|
||||
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
|
||||
|
||||
yield CheckpointTuple(
|
||||
config={
|
||||
list_config = cast(
|
||||
RunnableConfig,
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
)
|
||||
channel_values = self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
)
|
||||
self._resolve_delta_channels(list_config, channel_values)
|
||||
|
||||
yield CheckpointTuple(
|
||||
config=list_config,
|
||||
checkpoint={
|
||||
**checkpoint_,
|
||||
"channel_values": self._load_blobs(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_["channel_versions"],
|
||||
),
|
||||
"channel_values": channel_values,
|
||||
},
|
||||
metadata=metadata,
|
||||
parent_config=(
|
||||
|
||||
@@ -64,6 +64,14 @@ def _warn_once(
|
||||
logger.warning(msg, *args)
|
||||
|
||||
|
||||
def _get_delta_sentinel_cls() -> type:
|
||||
from langgraph.checkpoint.base import (
|
||||
DeltaChannelSentinel,
|
||||
) # lazy import avoids circular dep
|
||||
|
||||
return DeltaChannelSentinel
|
||||
|
||||
|
||||
class JsonPlusSerializer(SerializerProtocol):
|
||||
"""Serializer that uses ormsgpack, with optional fallbacks.
|
||||
|
||||
@@ -256,6 +264,8 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return "bytes", obj
|
||||
elif isinstance(obj, bytearray):
|
||||
return "bytearray", obj
|
||||
elif isinstance(obj, _get_delta_sentinel_cls()):
|
||||
return "delta", b""
|
||||
else:
|
||||
try:
|
||||
return "msgpack", _msgpack_enc(obj)
|
||||
@@ -278,6 +288,10 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
elif type_ == "delta":
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
return DeltaChannelSentinel()
|
||||
elif self.pickle_fallback and type_ == "pickle":
|
||||
return pickle.loads(data_)
|
||||
else:
|
||||
|
||||
@@ -997,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_channel_sentinel_serde_round_trip() -> None:
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
serde = JsonPlusSerializer()
|
||||
original = DeltaChannelSentinel()
|
||||
type_tag, blob = serde.dumps_typed(original)
|
||||
assert type_tag == "delta"
|
||||
loaded = serde.loads_typed((type_tag, blob))
|
||||
assert isinstance(loaded, DeltaChannelSentinel)
|
||||
|
||||
@@ -320,3 +320,74 @@ 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 DeltaChannelSentinel for delta channels (reconstruction deferred)."""
|
||||
from langgraph.checkpoint.base import (
|
||||
DeltaChannelSentinel,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
v1 = "00000000000000000000000000000001.0000000000000000"
|
||||
|
||||
sentinel = DeltaChannelSentinel()
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(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 isinstance(result[channel], DeltaChannelSentinel)
|
||||
|
||||
def test_get_channel_writes_collects_writes(self) -> None:
|
||||
"""get_channel_writes collects per-step writes oldest→newest."""
|
||||
from langgraph.checkpoint.base import empty_checkpoint
|
||||
|
||||
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"),
|
||||
}
|
||||
# cp1 has a write for channel
|
||||
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
||||
"task1",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "hi"}),
|
||||
"",
|
||||
)
|
||||
# cp2 has a write for channel
|
||||
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
|
||||
"task2",
|
||||
channel,
|
||||
serde.dumps_typed({"content": "bye"}),
|
||||
"",
|
||||
)
|
||||
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": "cp2",
|
||||
}
|
||||
}
|
||||
result = saver.get_channel_writes(config, channel)
|
||||
assert result == [{"content": "hi"}, {"content": "bye"}]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -119,3 +119,12 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
|
||||
Returns `True` if the channel was updated, `False` otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
|
||||
"""Called after checkpoint() with the assigned version, and after
|
||||
from_checkpoint() with the current channel version.
|
||||
|
||||
No-op by default. Override in channels that track their own version
|
||||
for incremental checkpointing (e.g. DeltaChannel).
|
||||
"""
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
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
|
||||
from langgraph.errors import EmptyChannelError
|
||||
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
class DeltaChannel(
|
||||
Generic[Value], BaseChannel[list[Value], Value, DeltaChannelSentinel]
|
||||
):
|
||||
"""A channel that stores only a sentinel in checkpoints; per-step writes are
|
||||
stored in checkpoint_writes and replayed through the operator at load time.
|
||||
|
||||
Use with append-style reducers (e.g. `add_messages`) on long-running threads
|
||||
to eliminate O(N²) blob growth — storage is O(N) using the writes table that
|
||||
every checkpointer already maintains.
|
||||
|
||||
Works with all checkpointers. Savers with dedicated implementations
|
||||
(InMemorySaver, PostgresSaver) reconstruct in one pass; others fall back to
|
||||
walking the checkpoint list.
|
||||
|
||||
Usage::
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
|
||||
# Dict-type reducer (type inferred from the Annotated outer type):
|
||||
files: Annotated[dict, DeltaChannel(merge_files)]
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "operator")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[list[Value], Any], list[Value]],
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.value: list[Value] = []
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
return False
|
||||
if (
|
||||
self.operator.__name__ != "<lambda>"
|
||||
and other.operator.__name__ != "<lambda>"
|
||||
):
|
||||
return self.operator is other.operator
|
||||
return True
|
||||
|
||||
@property
|
||||
def ValueType(self) -> Any:
|
||||
return list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
@property
|
||||
def UpdateType(self) -> Any:
|
||||
return self.typ | list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
new.value = self.value if self.value is MISSING else self.value.copy()
|
||||
return new
|
||||
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
new = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
try:
|
||||
new.value = new.typ()
|
||||
except Exception:
|
||||
new.value = []
|
||||
elif isinstance(checkpoint, list):
|
||||
# Flat list of write values (oldest→newest) from get_channel_writes.
|
||||
try:
|
||||
value: Any = new.typ()
|
||||
except Exception:
|
||||
value = []
|
||||
for write in checkpoint:
|
||||
value = new.operator(value, write)
|
||||
new.value = value
|
||||
else:
|
||||
# Backward compat: plain accumulated value (e.g. from a migrated thread).
|
||||
try:
|
||||
new.value = list(checkpoint)
|
||||
except Exception:
|
||||
new.value = []
|
||||
return new
|
||||
|
||||
def update(self, values: Sequence[Any]) -> bool:
|
||||
if not values:
|
||||
return False
|
||||
seen_overwrite = False
|
||||
for value in values:
|
||||
is_overwrite, overwrite_value = _get_overwrite(value)
|
||||
if is_overwrite:
|
||||
if seen_overwrite:
|
||||
from langgraph.errors import (
|
||||
ErrorCode,
|
||||
InvalidUpdateError,
|
||||
create_error_message,
|
||||
)
|
||||
|
||||
msg = create_error_message(
|
||||
message="Can receive only one Overwrite value per super-step.",
|
||||
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
|
||||
)
|
||||
raise InvalidUpdateError(msg)
|
||||
self.value = (
|
||||
list(overwrite_value) if overwrite_value is not None else self.typ()
|
||||
)
|
||||
seen_overwrite = True
|
||||
elif not seen_overwrite:
|
||||
base = self.typ() if self.value is MISSING else self.value
|
||||
self.value = self.operator(base, value)
|
||||
return True
|
||||
|
||||
def get(self) -> list[Value]:
|
||||
if self.value is MISSING:
|
||||
raise EmptyChannelError()
|
||||
return self.value
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> DeltaChannelSentinel:
|
||||
return DeltaChannelSentinel()
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import collections.abc
|
||||
import inspect
|
||||
import logging
|
||||
import typing
|
||||
@@ -47,7 +48,8 @@ from langgraph._internal._pydantic import create_model
|
||||
from langgraph._internal._runnable import coerce_to_runnable
|
||||
from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate, _strip_extras
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
|
||||
from langgraph.channels.named_barrier_value import (
|
||||
@@ -1082,6 +1084,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 +1670,18 @@ 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__"):
|
||||
outer = _strip_extras(typ.__origin__)
|
||||
if outer in (
|
||||
collections.abc.Sequence,
|
||||
collections.abc.MutableSequence,
|
||||
):
|
||||
outer = list
|
||||
item.typ = outer
|
||||
try:
|
||||
item.value = outer()
|
||||
except Exception:
|
||||
item.value = []
|
||||
return item
|
||||
elif isclass(item) and issubclass(item, BaseChannel):
|
||||
# ex, Annotated[int, EphemeralValue, SomeOtherAnnotation]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -12,6 +13,8 @@ from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
LATEST_VERSION = 4
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
return Checkpoint(
|
||||
@@ -67,13 +70,12 @@ def channels_from_checkpoint(
|
||||
channel_specs[k] = v
|
||||
else:
|
||||
managed_specs[k] = v
|
||||
return (
|
||||
{
|
||||
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
for k, v in channel_specs.items()
|
||||
},
|
||||
managed_specs,
|
||||
)
|
||||
channels: dict[str, BaseChannel] = {}
|
||||
for k, v in channel_specs.items():
|
||||
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
ch.after_checkpoint(checkpoint["channel_versions"].get(k), checkpoint.get("id"))
|
||||
channels[k] = ch
|
||||
return channels, managed_specs
|
||||
|
||||
|
||||
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
|
||||
@@ -891,6 +891,12 @@ class PregelLoop:
|
||||
id=self.checkpoint["id"] if exiting else None,
|
||||
updated_channels=self.updated_channels,
|
||||
)
|
||||
if do_checkpoint and self.channels:
|
||||
for k, ch in self.channels.items():
|
||||
ch.after_checkpoint(
|
||||
self.checkpoint["channel_versions"].get(k),
|
||||
self.checkpoint.get("id"),
|
||||
)
|
||||
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
|
||||
if TASKS in self.checkpoint["channel_values"] and any(
|
||||
isinstance(channel, UntrackedValue) for channel in self.channels.values()
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import functools
|
||||
import inspect
|
||||
import re
|
||||
import textwrap
|
||||
import types
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
@@ -64,46 +66,74 @@ def find_subgraph_pregel(candidate: Runnable) -> PregelProtocol | None:
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=256)
|
||||
def _get_nonlocal_names(code: types.CodeType) -> frozenset[str]:
|
||||
"""Return the set of nonlocal variable names referenced by a function.
|
||||
|
||||
Cached by code object so the expensive source fetch + AST parse only
|
||||
happens once per unique function definition across repeated graph compiles.
|
||||
|
||||
Args:
|
||||
code: The code object of the function to analyse.
|
||||
|
||||
Returns:
|
||||
Frozenset of variable names that the function reads from its enclosing
|
||||
scope (free variables and globals referenced in function bodies).
|
||||
"""
|
||||
try:
|
||||
source = inspect.getsource(code)
|
||||
tree = ast.parse(textwrap.dedent(source))
|
||||
visitor = FunctionNonLocals()
|
||||
visitor.visit(tree)
|
||||
return frozenset(visitor.nonlocals)
|
||||
except (SyntaxError, TypeError, OSError, SystemError):
|
||||
return frozenset()
|
||||
|
||||
|
||||
def get_function_nonlocals(func: Callable) -> list[Any]:
|
||||
"""Get the nonlocal variables accessed by a function.
|
||||
|
||||
The expensive source-parsing step is cached by code object; only the
|
||||
cheap closure-variable lookup runs on every call.
|
||||
|
||||
Args:
|
||||
func: The function to check.
|
||||
|
||||
Returns:
|
||||
List[Any]: The nonlocal variables accessed by the function.
|
||||
"""
|
||||
try:
|
||||
code = inspect.getsource(func)
|
||||
tree = ast.parse(textwrap.dedent(code))
|
||||
visitor = FunctionNonLocals()
|
||||
visitor.visit(tree)
|
||||
values: list[Any] = []
|
||||
closure = (
|
||||
inspect.getclosurevars(func.__wrapped__)
|
||||
if hasattr(func, "__wrapped__") and callable(func.__wrapped__)
|
||||
else inspect.getclosurevars(func)
|
||||
)
|
||||
candidates = {**closure.globals, **closure.nonlocals}
|
||||
for k, v in candidates.items():
|
||||
if k in visitor.nonlocals:
|
||||
values.append(v)
|
||||
for kk in visitor.nonlocals:
|
||||
if "." in kk and kk.startswith(k):
|
||||
vv = v
|
||||
for part in kk.split(".")[1:]:
|
||||
if vv is None:
|
||||
break
|
||||
else:
|
||||
try:
|
||||
vv = getattr(vv, part)
|
||||
except AttributeError:
|
||||
break
|
||||
else:
|
||||
values.append(vv)
|
||||
except (SyntaxError, TypeError, OSError, SystemError):
|
||||
actual_func = (
|
||||
func.__wrapped__
|
||||
if hasattr(func, "__wrapped__") and callable(func.__wrapped__)
|
||||
else func
|
||||
)
|
||||
# Fast path: no free variables means nothing to scan.
|
||||
if not actual_func.__code__.co_freevars:
|
||||
return []
|
||||
|
||||
nonlocal_names = _get_nonlocal_names(actual_func.__code__)
|
||||
if not nonlocal_names:
|
||||
return []
|
||||
|
||||
closure = inspect.getclosurevars(actual_func)
|
||||
candidates = {**closure.globals, **closure.nonlocals}
|
||||
values: list[Any] = []
|
||||
for k, v in candidates.items():
|
||||
if k in nonlocal_names:
|
||||
values.append(v)
|
||||
for kk in nonlocal_names:
|
||||
if "." in kk and kk.startswith(k):
|
||||
vv = v
|
||||
for part in kk.split(".")[1:]:
|
||||
if vv is None:
|
||||
break
|
||||
else:
|
||||
try:
|
||||
vv = getattr(vv, part)
|
||||
except AttributeError:
|
||||
break
|
||||
else:
|
||||
values.append(vv)
|
||||
return values
|
||||
|
||||
|
||||
|
||||
@@ -1049,13 +1049,14 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1168,13 +1169,14 @@ class Pregel(
|
||||
|
||||
step = saved.metadata.get("step", -1) + 1
|
||||
stop = step + 2
|
||||
checkpoint = saved.checkpoint
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
next_tasks = prepare_next_tasks(
|
||||
saved.checkpoint,
|
||||
checkpoint,
|
||||
saved.pending_writes or [],
|
||||
self.nodes,
|
||||
channels,
|
||||
@@ -1520,9 +1522,8 @@ class Pregel(
|
||||
saved = checkpointer.get_tuple(config)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
@@ -1966,9 +1967,8 @@ class Pregel(
|
||||
saved = await checkpointer.aget_tuple(config)
|
||||
if saved is not None:
|
||||
self._migrate_checkpoint(saved.checkpoint)
|
||||
checkpoint = (
|
||||
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
|
||||
)
|
||||
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
|
||||
checkpoint_previous_versions = (
|
||||
saved.checkpoint["channel_versions"].copy() if saved else {}
|
||||
)
|
||||
|
||||
@@ -117,3 +117,291 @@ def test_untracked_value() -> None:
|
||||
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
|
||||
with pytest.raises(EmptyChannelError):
|
||||
new_channel.get()
|
||||
|
||||
|
||||
def test_delta_channel_basic_two_steps() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
|
||||
# Step 1: one message added
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaChannelSentinel)
|
||||
|
||||
# Step 2: another message
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert isinstance(d2, DeltaChannelSentinel)
|
||||
|
||||
# Full accumulated value is preserved in memory
|
||||
assert len(ch.get()) == 2
|
||||
assert ch.get()[0].content == "hi"
|
||||
assert ch.get()[1].content == "hello"
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_writes_list() -> None:
|
||||
"""from_checkpoint with a flat list of individual writes replays them through the operator."""
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
# Each element is one write value (as stored in checkpoint_writes)
|
||||
writes = [
|
||||
HumanMessage(content="hi", id="h1"),
|
||||
AIMessage(content="hello", id="a1"),
|
||||
HumanMessage(content="bye", id="h2"),
|
||||
]
|
||||
ch = spec.from_checkpoint(writes)
|
||||
msgs = ch.get()
|
||||
assert len(msgs) == 3
|
||||
assert msgs[0].content == "hi"
|
||||
assert msgs[1].content == "hello"
|
||||
assert msgs[2].content == "bye"
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat
|
||||
spec = DeltaChannel(add_messages)
|
||||
old_value = [HumanMessage(content="old", id="h1")]
|
||||
ch = spec.from_checkpoint(old_value)
|
||||
assert ch.get() == old_value
|
||||
|
||||
|
||||
def test_delta_channel_overwrite() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
ch.update([HumanMessage(content="old", id="h1")])
|
||||
|
||||
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
|
||||
d = ch.checkpoint()
|
||||
assert isinstance(d, DeltaChannelSentinel)
|
||||
# After overwrite, value is reset to only the new message
|
||||
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."""
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
|
||||
# Step 1: add two messages
|
||||
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"),
|
||||
]
|
||||
|
||||
# Step 2: remove the AI message
|
||||
ch.update([RemoveMessage(id="a1")])
|
||||
assert ch.get() == [HumanMessage(content="hi", id="h1")]
|
||||
|
||||
# Replay the writes list from scratch — must reproduce the post-remove state
|
||||
writes = [
|
||||
HumanMessage(content="hi", id="h1"),
|
||||
AIMessage(content="hello", id="a1"),
|
||||
RemoveMessage(id="a1"),
|
||||
]
|
||||
ch2 = spec.from_checkpoint(writes)
|
||||
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."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
|
||||
# Step 1: add a message
|
||||
ch.update([HumanMessage(content="original", id="h1")])
|
||||
|
||||
# Step 2: update the same message by ID
|
||||
ch.update([HumanMessage(content="updated", id="h1")])
|
||||
assert ch.get() == [HumanMessage(content="updated", id="h1")]
|
||||
|
||||
# Replay writes — must produce the updated message, not the original
|
||||
writes = [
|
||||
HumanMessage(content="original", id="h1"),
|
||||
HumanMessage(content="updated", id="h1"),
|
||||
]
|
||||
ch2 = spec.from_checkpoint(writes)
|
||||
assert len(ch2.get()) == 1
|
||||
assert ch2.get()[0].content == "updated"
|
||||
|
||||
|
||||
def test_delta_channel_checkpoint_returns_sentinel() -> None:
|
||||
"""checkpoint() always returns DeltaChannelSentinel regardless of state."""
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
assert isinstance(ch.checkpoint(), DeltaChannelSentinel)
|
||||
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
assert isinstance(ch.checkpoint(), DeltaChannelSentinel)
|
||||
|
||||
|
||||
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
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)
|
||||
|
||||
# get_tuple must return a resolved list (not DeltaChannelSentinel)
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
saved = saver.get_tuple(config)
|
||||
assert saved is not None
|
||||
assert "messages" in saved.checkpoint["channel_values"]
|
||||
assert not isinstance(
|
||||
saved.checkpoint["channel_values"]["messages"], DeltaChannelSentinel
|
||||
)
|
||||
assert isinstance(saved.checkpoint["channel_values"]["messages"], list)
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
|
||||
|
||||
|
||||
def _delta_channel_with_type(operator, typ):
|
||||
"""Build a DeltaChannel with an explicit type via the Annotated injection path."""
|
||||
from typing import Annotated
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.state import _get_channel
|
||||
|
||||
return _get_channel("_test", Annotated[typ, DeltaChannel(operator)])
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_fresh_channel() -> None:
|
||||
"""DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint."""
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
# Should be available (not raise EmptyChannelError) and start empty
|
||||
assert ch.is_available()
|
||||
assert ch.get() == {}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_basic_updates() -> None:
|
||||
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
|
||||
|
||||
ch.update([{"a": 1}])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaChannelSentinel)
|
||||
|
||||
ch.update([{"b": 2}])
|
||||
d2 = ch.checkpoint()
|
||||
assert isinstance(d2, DeltaChannelSentinel)
|
||||
|
||||
assert ch.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_writes_reconstruction() -> None:
|
||||
"""from_checkpoint with a writes list replays correctly through a dict merge reducer."""
|
||||
|
||||
def merge_dicts(left: dict, right: dict) -> dict:
|
||||
return {**left, **right}
|
||||
|
||||
spec = _delta_channel_with_type(merge_dicts, dict)
|
||||
# Each element is one write value (oldest→newest)
|
||||
writes = [{"a": 1}, {"b": 2}, {"c": 3}]
|
||||
ch = spec.from_checkpoint(writes)
|
||||
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 (deepagents pattern)."""
|
||||
|
||||
def merge_files(left: dict | None, right: dict) -> dict:
|
||||
if left is None:
|
||||
return {k: v for k, v in right.items() if v is not None}
|
||||
result = {**left}
|
||||
for k, v in right.items():
|
||||
if v is None:
|
||||
result.pop(k, None)
|
||||
else:
|
||||
result[k] = v
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING)
|
||||
|
||||
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
|
||||
|
||||
# Delete file1, add file3
|
||||
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
||||
|
||||
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
# Confirm writes reconstruction produces the same result
|
||||
writes = [
|
||||
{"file1.py": "content1", "file2.py": "content2"},
|
||||
{"file1.py": None, "file3.py": "content3"},
|
||||
]
|
||||
spec = _delta_channel_with_type(merge_files, dict)
|
||||
ch2 = spec.from_checkpoint(writes)
|
||||
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
"""Benchmark: DeltaChannel vs BinaryOperatorAggregate storage and time.
|
||||
|
||||
Run directly: python tests/test_delta_channel_benchmark.py
|
||||
Run via pytest: pytest tests/test_delta_channel_benchmark.py -s
|
||||
|
||||
Simulates realistic multi-turn conversations with paragraph-length messages
|
||||
(~100 tokens each) scaling up to 1M-token-equivalent histories.
|
||||
|
||||
Token estimates: 1 token ≈ 4 chars; each turn ≈ 200 tokens (human + AI).
|
||||
A 1M-token conversation ≈ 5,000 turns of realistic messages.
|
||||
|
||||
DeltaChannel stores only a zero-byte sentinel in checkpoint_blobs; the actual
|
||||
write data lives in checkpoint_writes (already stored there). Reconstruction
|
||||
walks the parent chain and replays writes through the operator — O(N) total
|
||||
storage vs O(N²) for plain add_messages.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import 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 add_messages
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
|
||||
_SQLITE_AVAILABLE = True
|
||||
except ImportError:
|
||||
_SQLITE_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
_POSTGRES_AVAILABLE = True
|
||||
_POSTGRES_URI = (
|
||||
"postgres://postgres:postgres@localhost:5441/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."
|
||||
)
|
||||
|
||||
_AI_TEMPLATE = (
|
||||
"Great question about {topic}. The key insight here is that {concern} introduces "
|
||||
"a subtle ordering dependency that most teams overlook until they hit it in production. "
|
||||
"For your {component} layer specifically, I'd recommend starting with a careful audit "
|
||||
"of the interface boundaries before making any structural changes. This will give you "
|
||||
"a clear picture of the blast radius and let you sequence the migration safely."
|
||||
)
|
||||
|
||||
_TOPICS = [
|
||||
"distributed tracing",
|
||||
"eventual consistency",
|
||||
"schema migration",
|
||||
"backpressure handling",
|
||||
"idempotency guarantees",
|
||||
"cache invalidation",
|
||||
"connection pooling",
|
||||
"rate limiting",
|
||||
"circuit breaking",
|
||||
"observability pipelines",
|
||||
]
|
||||
|
||||
_CONCERNS = [
|
||||
"concurrency model",
|
||||
"retry semantics",
|
||||
"state management",
|
||||
"error propagation",
|
||||
"latency budget",
|
||||
]
|
||||
|
||||
_COMPONENTS = [
|
||||
"persistence",
|
||||
"routing",
|
||||
"ingestion",
|
||||
"aggregation",
|
||||
"serialization",
|
||||
]
|
||||
|
||||
|
||||
def _human_content(i: int) -> str:
|
||||
return _HUMAN_TEMPLATE.format(
|
||||
topic=_TOPICS[i % len(_TOPICS)],
|
||||
concern=_CONCERNS[i % len(_CONCERNS)],
|
||||
component=_COMPONENTS[i % len(_COMPONENTS)],
|
||||
)
|
||||
|
||||
|
||||
def _ai_content(i: int) -> str:
|
||||
return _AI_TEMPLATE.format(
|
||||
topic=_TOPICS[i % len(_TOPICS)],
|
||||
concern=_CONCERNS[i % len(_CONCERNS)],
|
||||
component=_COMPONENTS[i % len(_COMPONENTS)],
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# State definitions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BinaryState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
|
||||
|
||||
class DeltaState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_graph(state_cls: type, checkpointer: Any = None) -> Any:
|
||||
def human_node(state: Any) -> dict:
|
||||
return {}
|
||||
|
||||
def ai_node(state: Any) -> dict:
|
||||
i = len(state["messages"]) // 2
|
||||
return {"messages": [AIMessage(content=_ai_content(i), id=f"a{i}")]}
|
||||
|
||||
g = StateGraph(state_cls)
|
||||
g.add_node("human", human_node)
|
||||
g.add_node("ai", ai_node)
|
||||
g.add_edge("human", "ai")
|
||||
g.add_edge("ai", END)
|
||||
g.set_entry_point("human")
|
||||
return g.compile(checkpointer=checkpointer or MemorySaver())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Measurement helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _total_blob_bytes(saver: MemorySaver) -> int:
|
||||
total = 0
|
||||
for (_, _, _, _), (type_tag, blob) in saver.blobs.items():
|
||||
if blob is not None:
|
||||
total += len(blob)
|
||||
return total
|
||||
|
||||
|
||||
def _run_turns(
|
||||
n_turns: int,
|
||||
state_cls: type,
|
||||
checkpointer: Any = None,
|
||||
) -> tuple[float, float, int]:
|
||||
"""Run n_turns conversation turns.
|
||||
|
||||
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
|
||||
blob_bytes is -1 for savers without in-memory blob stores (e.g. SQLite).
|
||||
Read latency is measured as the time to invoke the graph with no new
|
||||
messages after the full history is built — this forces state rehydration.
|
||||
"""
|
||||
graph = _make_graph(state_cls, checkpointer)
|
||||
config = {"configurable": {"thread_id": "bench"}}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
for i in range(n_turns):
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
|
||||
config,
|
||||
)
|
||||
write_elapsed = time.perf_counter() - t0
|
||||
|
||||
# Measure read/rehydration: get_state forces the channel to rebuild
|
||||
t1 = time.perf_counter()
|
||||
for _ in range(5):
|
||||
graph.get_state(config)
|
||||
read_elapsed = (time.perf_counter() - t1) / 5
|
||||
|
||||
if isinstance(graph.checkpointer, MemorySaver):
|
||||
blob_bytes = _total_blob_bytes(graph.checkpointer)
|
||||
else:
|
||||
blob_bytes = -1
|
||||
return write_elapsed, read_elapsed, blob_bytes
|
||||
|
||||
|
||||
def _fmt_bytes(n: int) -> str:
|
||||
if n >= 1_000_000:
|
||||
return f"{n / 1_000_000:.1f} MB"
|
||||
if n >= 1_000:
|
||||
return f"{n / 1_000:.1f} KB"
|
||||
return f"{n} B"
|
||||
|
||||
|
||||
def _approx_tokens(n_turns: int) -> str:
|
||||
# ~100 tokens human + ~100 tokens AI per turn
|
||||
tokens = n_turns * 200
|
||||
if tokens >= 1_000_000:
|
||||
return f"~{tokens / 1_000_000:.1f}M tok"
|
||||
if tokens >= 1_000:
|
||||
return f"~{tokens / 1_000:.0f}K tok"
|
||||
return f"~{tokens} tok"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark matrix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Turn counts chosen to demonstrate O(N²) vs O(N) storage growth without running too long.
|
||||
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
|
||||
TURN_COUNTS = [10, 25, 50, 100, 500]
|
||||
|
||||
|
||||
def _checkpointer_factories() -> list[tuple[str, Any]]:
|
||||
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
|
||||
return [("InMemory", None)]
|
||||
|
||||
|
||||
def run_benchmark() -> None:
|
||||
print()
|
||||
print(
|
||||
"DeltaChannel vs add_messages (BinaryOperatorAggregate) — checkpoint storage & latency"
|
||||
)
|
||||
print("Simulating realistic multi-turn conversations up to ~1M-token histories")
|
||||
print("(5,000 turns × ~200 tokens/turn ≈ 1M tokens — Claude's full context window)")
|
||||
print()
|
||||
|
||||
checkpointers: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
if _POSTGRES_AVAILABLE:
|
||||
try:
|
||||
import psycopg
|
||||
|
||||
psycopg.connect(_POSTGRES_URI).close()
|
||||
checkpointers.append(("Postgres (recursive CTE)", "postgres"))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for cp_label, cp_hint in checkpointers:
|
||||
print(f"--- Checkpointer: {cp_label} ---")
|
||||
_run_benchmark_for_checkpointer(cp_hint)
|
||||
|
||||
|
||||
def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
import contextlib
|
||||
import tempfile
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _make_saver():
|
||||
if cp_hint is None:
|
||||
yield None
|
||||
elif cp_hint == "postgres":
|
||||
with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver:
|
||||
saver.setup()
|
||||
with saver._cursor() as cur:
|
||||
cur.execute("DELETE FROM checkpoints WHERE thread_id = 'bench'")
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoint_blobs WHERE thread_id = 'bench'"
|
||||
)
|
||||
cur.execute(
|
||||
"DELETE FROM checkpoint_writes WHERE thread_id = 'bench'"
|
||||
)
|
||||
yield saver
|
||||
else:
|
||||
with tempfile.NamedTemporaryFile(suffix=".db") as f:
|
||||
with SqliteSaver.from_conn_string(f.name) as saver:
|
||||
yield saver
|
||||
|
||||
rows = []
|
||||
for turns in TURN_COUNTS:
|
||||
with _make_saver() as saver:
|
||||
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt))
|
||||
|
||||
# ── Table 1: Storage ─────────────────────────────────────────────────────
|
||||
W = 70
|
||||
print("Storage (checkpoint blob bytes)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'savings':>8}"
|
||||
)
|
||||
print("-" * W)
|
||||
storage_results = []
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt in rows:
|
||||
if b_bytes < 0:
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>8}"
|
||||
)
|
||||
else:
|
||||
ratio = b_bytes / d_bytes if d_bytes else float("inf")
|
||||
storage_results.append((turns, b_bytes, d_bytes, ratio))
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} "
|
||||
f"{ratio:>7.0f}x"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
# ── Table 2: Read latency ─────────────────────────────────────────────────
|
||||
print("Read latency (avg of 5 get_state calls)")
|
||||
print("=" * W)
|
||||
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt in rows:
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} "
|
||||
f"{b_rt * 1000:>10.1f}ms {d_rt * 1000:>10.1f}ms"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
if storage_results:
|
||||
turns, b_bytes, d_bytes, ratio = storage_results[-1]
|
||||
b_rt = rows[-1][-2]
|
||||
d_rt = rows[-1][-1]
|
||||
print(
|
||||
f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(d_bytes)} ({ratio:.0f}x less storage); "
|
||||
f"read {b_rt * 1000:.1f}ms → {d_rt * 1000:.1f}ms"
|
||||
)
|
||||
print()
|
||||
|
||||
print("Legend:")
|
||||
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
|
||||
print(
|
||||
" delta = DeltaChannel(add_messages) — O(N) storage, full chain replay"
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pytest entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
|
||||
)
|
||||
def test_delta_channel_benchmark(capsys: Any) -> None:
|
||||
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
|
||||
with capsys.disabled():
|
||||
run_benchmark()
|
||||
|
||||
# Correctness assertion: DeltaChannel must use less storage at scale.
|
||||
for turns in [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}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Script entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_benchmark()
|
||||
sys.exit(0)
|
||||
@@ -9400,3 +9400,190 @@ def test_fork_does_not_apply_pending_writes(
|
||||
|
||||
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
|
||||
assert result == {"value": 121}
|
||||
|
||||
|
||||
async def test_delta_channel_end_to_end_inmemory() -> None:
|
||||
"""Full graph run: DeltaChannel accumulates correctly across multiple turns."""
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
n = len(state["messages"])
|
||||
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-test-1"}}
|
||||
|
||||
# Turn 1
|
||||
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
||||
# Turn 2
|
||||
graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config)
|
||||
# Turn 3
|
||||
graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# 3 human + 3 AI = 6 total
|
||||
assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].content == "hello"
|
||||
assert msgs[2].content == "world"
|
||||
assert msgs[4].content == "bye"
|
||||
assert msgs[1].content == "reply-1"
|
||||
assert msgs[3].content == "reply-3"
|
||||
assert msgs[5].content == "reply-5"
|
||||
|
||||
|
||||
async def test_delta_channel_time_travel() -> None:
|
||||
"""Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas."""
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
counter["n"] += 1
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
|
||||
]
|
||||
}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
saver = InMemorySaver()
|
||||
graph = builder.compile(checkpointer=saver)
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-time-travel"}}
|
||||
|
||||
# Run 2 turns: h1→ai-1, h2→ai-2
|
||||
graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config)
|
||||
graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config)
|
||||
|
||||
# Find the checkpoint after turn 1 (2 messages: h1 + ai-1)
|
||||
history = list(graph.get_state_history(config))
|
||||
after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2)
|
||||
|
||||
assert len(after_turn1.values["messages"]) == 2
|
||||
assert after_turn1.values["messages"][0].content == "h1"
|
||||
assert after_turn1.values["messages"][1].content == "ai-1"
|
||||
|
||||
# Resume from turn-1 checkpoint: inject h3, expect 3 messages total (h1, ai-1, ai-N)
|
||||
# NOT 5 messages (turn-2 deltas must not bleed into the resumed run)
|
||||
result = graph.invoke(
|
||||
{"messages": [HumanMessage(content="h3", id="h3")]},
|
||||
after_turn1.config,
|
||||
)
|
||||
msgs = result["messages"]
|
||||
# Should be: h1, ai-1, h3, ai-N — 4 messages total
|
||||
assert len(msgs) == 4, (
|
||||
f"expected 4 messages after time-travel resume, got {len(msgs)}: {msgs}"
|
||||
)
|
||||
assert msgs[0].content == "h1"
|
||||
assert msgs[1].content == "ai-1"
|
||||
assert msgs[2].content == "h3"
|
||||
|
||||
|
||||
async def test_delta_channel_remove_message_end_to_end() -> None:
|
||||
"""RemoveMessage inside a DeltaChannel graph must persist and reload correctly."""
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
return {"messages": [AIMessage(content="reply", id="ai-1")]}
|
||||
|
||||
def delete_first(state: State) -> dict:
|
||||
# removes the first message
|
||||
return {"messages": [RemoveMessage(id=state["messages"][0].id)]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_node("delete_first", delete_first)
|
||||
builder.add_edge(START, "respond")
|
||||
builder.add_edge("respond", "delete_first")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-remove-test"}}
|
||||
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# h1 was removed, only ai-1 should remain
|
||||
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].id == "ai-1"
|
||||
|
||||
# A subsequent turn must reconstruct from the checkpoint correctly
|
||||
graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config)
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed
|
||||
# more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2
|
||||
# just assert h1 is still gone
|
||||
assert all(m.id != "h1" for m in msgs), (
|
||||
"h1 should still be absent after second turn"
|
||||
)
|
||||
|
||||
|
||||
async def test_delta_channel_update_by_id_end_to_end() -> None:
|
||||
"""Updating a message by ID via DeltaChannel must persist and reload correctly."""
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
def update_msg(state: State) -> dict:
|
||||
# re-send h1 with updated content
|
||||
return {"messages": [HumanMessage(content="updated", id="h1")]}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("update_msg", update_msg)
|
||||
builder.add_edge(START, "update_msg")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
config = {"configurable": {"thread_id": "diff-update-id-test"}}
|
||||
graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
|
||||
assert msgs[0].content == "updated"
|
||||
assert msgs[0].id == "h1"
|
||||
|
||||
# Second turn: verify the updated state is the base for further accumulation
|
||||
graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config)
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
ids = [m.id for m in msgs]
|
||||
assert "h1" in ids # h1 persists (updated, not duplicated)
|
||||
assert "h2" in ids
|
||||
assert ids.count("h1") == 1, "h1 must not be duplicated"
|
||||
|
||||
@@ -427,3 +427,49 @@ def test_callback_manager_copies_configurable_ids_to_tracing_metadata() -> None:
|
||||
"thread_id": "th-123",
|
||||
"user_id": "uid-1",
|
||||
}
|
||||
|
||||
|
||||
def test_get_nonlocal_names_cached_by_code_object() -> None:
|
||||
"""_get_nonlocal_names caches by code object so repeated calls are cheap."""
|
||||
from langgraph.pregel._utils import _get_nonlocal_names
|
||||
|
||||
x = 1
|
||||
|
||||
def my_func() -> int:
|
||||
return x
|
||||
|
||||
result1 = _get_nonlocal_names(my_func.__code__)
|
||||
result2 = _get_nonlocal_names(my_func.__code__)
|
||||
|
||||
# Same frozenset instance returned (cache hit)
|
||||
assert result1 is result2
|
||||
assert "x" in result1
|
||||
|
||||
|
||||
def test_get_function_nonlocals_fast_path_no_freevars() -> None:
|
||||
"""Functions with no free variables return [] without AST parsing."""
|
||||
from langgraph.pregel._utils import _get_nonlocal_names, get_function_nonlocals
|
||||
|
||||
cache_info_before = _get_nonlocal_names.cache_info()
|
||||
|
||||
def pure_func(a: int, b: int) -> int:
|
||||
return a + b
|
||||
|
||||
result = get_function_nonlocals(pure_func)
|
||||
|
||||
# Should have returned early without touching the cache
|
||||
assert result == []
|
||||
assert _get_nonlocal_names.cache_info().misses == cache_info_before.misses
|
||||
|
||||
|
||||
def test_get_function_nonlocals_returns_closure_values() -> None:
|
||||
"""get_function_nonlocals correctly extracts values from closures."""
|
||||
from langgraph.pregel._utils import get_function_nonlocals
|
||||
|
||||
sentinel = object()
|
||||
|
||||
def my_func() -> object:
|
||||
return sentinel
|
||||
|
||||
result = get_function_nonlocals(my_func)
|
||||
assert sentinel in result
|
||||
|
||||
@@ -1842,9 +1842,17 @@ def _get_injection_from_type(
|
||||
return None
|
||||
|
||||
|
||||
# Cache keyed by tool object identity. Stores (tool, result) to keep a strong
|
||||
# reference that prevents GC from reusing the id for a different object.
|
||||
_INJECTED_ARGS_CACHE: dict[int, tuple[BaseTool, _InjectedArgs]] = {}
|
||||
|
||||
|
||||
def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
"""Extract all injected arguments from tool in a single pass.
|
||||
|
||||
Results are cached by tool identity so the expensive type-hint and schema
|
||||
inspection only runs once per unique tool object across ToolNode instances.
|
||||
|
||||
This function analyzes both the tool's input schema and function signature
|
||||
to identify all arguments that should be injected (state, store, runtime).
|
||||
|
||||
@@ -1854,6 +1862,11 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
Returns:
|
||||
_InjectedArgs structure containing all detected injections.
|
||||
"""
|
||||
tool_id = id(tool)
|
||||
entry = _INJECTED_ARGS_CACHE.get(tool_id)
|
||||
if entry is not None and entry[0] is tool:
|
||||
return entry[1]
|
||||
|
||||
# Get annotations from both schema and function signature
|
||||
full_schema = tool.get_input_schema()
|
||||
schema_annotations = get_all_basemodel_annotations(full_schema)
|
||||
@@ -1899,10 +1912,12 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
|
||||
if _get_injection_from_type(type_, ToolRuntime):
|
||||
runtime_arg = name
|
||||
|
||||
return _InjectedArgs(
|
||||
result = _InjectedArgs(
|
||||
state=state_args,
|
||||
store=store_arg,
|
||||
runtime=runtime_arg,
|
||||
all_injected_keys=all_injected_keys,
|
||||
_optional_state_args=_optional_state_args,
|
||||
)
|
||||
_INJECTED_ARGS_CACHE[tool_id] = (tool, result)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
# feat(channels): DeltaChannel — O(N) incremental checkpoint storage
|
||||
|
||||
## The problem
|
||||
|
||||
LangGraph checkpoints store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, that means each checkpoint blob contains the entire conversation history up to that point.
|
||||
|
||||
Storage cost grows **O(N²)** in the number of turns:
|
||||
|
||||
| Step | Checkpoint blob |
|
||||
|------|----------------|
|
||||
| 1 | [msg_1] |
|
||||
| 2 | [msg_1, msg_2] |
|
||||
| N | [msg_1, ..., msg_N] |
|
||||
|
||||
At 100K tokens of conversation data, a single thread accumulates ~250 MB; with large messages or file attachments costs scale even faster.
|
||||
|
||||
## The fix: `DeltaChannel`
|
||||
|
||||
`DeltaChannel` is an opt-in wrapper around any binary reducer that stores only a **sentinel marker** in `checkpoint_blobs` rather than the full accumulated value. The actual per-step writes stay in `checkpoint_writes` (which every checkpointer already writes unconditionally). At read time the saver walks the ancestor chain, collects all writes for the channel, and replays them through the reducer.
|
||||
|
||||
Storage scales **O(N)** — the sentinel blob is effectively zero bytes, and the writes table already exists.
|
||||
|
||||
```python
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
# Before: O(N²) storage
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
|
||||
# After: O(N) storage
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
|
||||
```
|
||||
|
||||
## Benchmarks
|
||||
|
||||
Simulated with realistic paragraph-length messages (~100 tokens each, ~400 chars). Each turn = one human + one AI message (~200 tokens total).
|
||||
|
||||
### Storage (InMemorySaver)
|
||||
|
||||
| turns | ctx | add_msgs | delta | savings |
|
||||
|------:|----:|---------:|------:|--------:|
|
||||
| 10 | ~2K tok | 108.6 KB | 4.0 KB | 27x |
|
||||
| 25 | ~5K tok | 649.0 KB | 10.1 KB | 64x |
|
||||
| 50 | ~10K tok | 2.6 MB | 20.2 KB | 126x |
|
||||
| 100 | ~20K tok | 10.2 MB | 40.5 KB | 251x |
|
||||
| 500 | ~100K tok | 252.6 MB | 202.8 KB | 1245x |
|
||||
|
||||
Savings grow with N because `add_messages` is O(N²) while `DeltaChannel` is O(N). The sentinel blob itself is essentially zero bytes.
|
||||
|
||||
### Read latency (avg of 5 `get_state` calls = cost per `invoke`)
|
||||
|
||||
| turns | ctx | add_msgs | delta |
|
||||
|------:|----:|---------:|------:|
|
||||
| 10 | ~2K tok | 0.1ms | 0.2ms |
|
||||
| 25 | ~5K tok | 0.2ms | 0.5ms |
|
||||
| 50 | ~10K tok | 0.4ms | 1.5ms |
|
||||
| 100 | ~20K tok | 0.7ms | 4.8ms |
|
||||
| 500 | ~100K tok | 5.8ms | 114.9ms |
|
||||
|
||||
**This cost is paid once per `invoke`/`stream` call, not per node.** Within a single invocation, all channels are loaded into memory once at the start and shared across every node — there is no per-node reconstruction. The 114.9ms at 500 turns is what you pay each time a user sends a new message, not on each step of the graph.
|
||||
|
||||
## How it works
|
||||
|
||||
**Write:** `DeltaChannel.checkpoint()` always emits `DeltaChannelSentinel()` — a tiny marker (zero payload bytes) stored in `checkpoint_blobs`. Per-step writes flow into `checkpoint_writes` as they normally do for every channel.
|
||||
|
||||
**Read:** The saver detects `DeltaChannelSentinel` values in `channel_values` and replaces them by calling `get_channel_writes` / `aget_channel_writes`, which walks the ancestor checkpoint chain and collects all writes for that channel (oldest→newest). `DeltaChannel.from_checkpoint()` replays those writes through the operator to reconstruct the full value.
|
||||
|
||||
**Saver implementations:**
|
||||
- `InMemorySaver` — direct dict traversal of `self.storage` and `self.writes`, no I/O
|
||||
- `PostgresSaver` (sync + async) — two queries: one cheap ID walk across the thread, one `ANY()` fetch of writes; no recursive CTE
|
||||
- All other savers — `BaseCheckpointSaver.get_channel_writes` fallback via `list()`, with a re-entrancy guard to prevent infinite recursion
|
||||
|
||||
## Changes
|
||||
|
||||
**`libs/checkpoint`**
|
||||
- `base/__init__.py` — add `DeltaChannelSentinel` marker dataclass; add `get_channel_writes` / `aget_channel_writes` to `BaseCheckpointSaver` with a `list()`-based fallback and re-entrancy guard
|
||||
|
||||
**`libs/checkpoint/memory`**
|
||||
- `memory/__init__.py` — `get_channel_writes` via direct dict traversal; `_resolve_delta_channels` helper called in `get_tuple` / `aget_tuple` to replace sentinels with reconstructed write lists
|
||||
|
||||
**`libs/langgraph`**
|
||||
- `channels/delta.py` — `DeltaChannel` implementation: `checkpoint()` always emits sentinel, `from_checkpoint()` replays writes list
|
||||
- `channels/__init__.py` — export `DeltaChannel`
|
||||
- `graph/state.py` — recognize `DeltaChannel` as a valid channel annotation
|
||||
- `pregel/_checkpoint.py` / `pregel/_loop.py` — wire `after_checkpoint` hook; call it after each checkpointing step so `DeltaChannel` can advance internal state
|
||||
|
||||
**`libs/checkpoint-postgres`**
|
||||
- `postgres/base.py` — `_get_channel_writes_cur` two-query ancestor walk (sync); `_resolve_delta_channels` called after `_load_blobs`
|
||||
- `postgres/aio.py` — `_aget_channel_writes_cur` (async counterpart)
|
||||
|
||||
## Open questions
|
||||
|
||||
**Should we add a compile-time capability check?**
|
||||
|
||||
Currently misconfiguring `DeltaChannel` with an unsupported saver only errors at runtime on first reload. A protocol-based check at `compile()` time would give an early warning without requiring a manual boolean flag.
|
||||
|
||||
**`snapshot_every` for bounded reconstruction cost?**
|
||||
|
||||
Both per-invoke read latency and total write wall time grow O(N) per invoke / O(N²) total as the conversation lengthens. A `snapshot_every` parameter — periodically store a full snapshot in `checkpoint_blobs` to cap chain depth — would bound reconstruction cost and is a natural follow-up once the core design is stable.
|
||||
|
||||
## Backwards compatibility
|
||||
|
||||
| Scenario | Behaviour |
|
||||
|----------|-----------|
|
||||
| Existing graph using `add_messages` | Unaffected — no code or schema changes |
|
||||
| `DeltaChannel` loading an old full-list checkpoint blob | Handled via backwards-compat path in `from_checkpoint` |
|
||||
| `DeltaChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported |
|
||||
| Time-travel to a past checkpoint | Ancestor walk uses the version at that checkpoint — correct by construction |
|
||||
| `Overwrite` value | Resets the effective chain; reconstruction starts from that step |
|
||||
|
||||
## Test plan
|
||||
|
||||
- [x] `DeltaChannel` unit tests: `update` → `checkpoint` lifecycle, `from_checkpoint` chain replay, backwards-compat with plain list, `Overwrite` resets chain
|
||||
- [x] `InMemorySaver` `get_channel_writes`: assembles write list from dict storage
|
||||
- [x] Serde round-trip for `DeltaChannelSentinel`
|
||||
- [x] End-to-end graph tests: multi-turn conversations accumulate correctly, time-travel reconstructs correct partial history
|
||||
- [x] `PostgresSaver` two-query chain reconstruction (sync + async)
|
||||
- [x] `BaseCheckpointSaver` fallback path via `list()` with re-entrancy guard
|
||||
- [x] Storage benchmark: `DeltaChannel` uses strictly less storage than `add_messages` at all measured turn counts
|
||||
|
||||
---
|
||||
|
||||
## Changes from previous base branch
|
||||
|
||||
The previous version stored `DeltaValue` objects (containing the per-step writes) directly in `checkpoint_blobs` and used a `DeltaChainValue` to represent the assembled chain. Reconstruction required a dedicated `get_delta_chain` / `aget_delta_chain` protocol and a recursive CTE in Postgres.
|
||||
|
||||
This version pivots to a simpler design:
|
||||
- **Sentinel in blobs, writes in `checkpoint_writes`** — `checkpoint_blobs` stores only a zero-byte `DeltaChannelSentinel` marker. The actual per-step data already lives in `checkpoint_writes` (written unconditionally by every checkpointer), so blob storage is essentially free. This is why storage savings jump to 1245x at 500 turns.
|
||||
- **No custom serde type for the delta payload** — `DeltaValue` / `DeltaChainValue` and the `"delta"` serde type tag are gone. Writes are deserialized with the same serde path they were originally written with.
|
||||
- **Postgres: two queries instead of a recursive CTE** — fetch all `(checkpoint_id, parent_checkpoint_id)` pairs for the thread, walk the ancestor chain in Python, then fetch writes with a plain `ANY()` filter.
|
||||
- **Universal fallback on `BaseCheckpointSaver`** — the base class now provides `get_channel_writes` via `list()`, so any third-party saver works without modification.
|
||||
- **`snapshot_every` removed** — deferred as a follow-up; the simpler design is easier to reason about and delivers larger storage savings.
|
||||
Reference in New Issue
Block a user