mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-30 19:59:40 +02:00
feat(delta-channel): store sentinel in blobs, reconstruct from checkpoint_writes
DeltaChannel.checkpoint() now returns a zero-byte DeltaChannelSentinel instead of duplicating delta data in checkpoint_blobs. Reconstruction walks the parent checkpoint chain via checkpoint_writes (which already holds per-step writes) and replays them through the operator. In-memory benchmark (100 turns, ~20K tokens): storage: 10.2 MB → 40.5 KB (251x reduction) read: 0.6ms → 7.9ms (reconstruction cost, amortized by storage savings) InMemorySaver and PostgresSaver override get_channel_writes() with efficient implementations (Python dict walk and recursive CTE respectively). The base class fallback uses self.list() with a thread-local recursion guard.
This commit is contained in:
@@ -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 (
|
||||
@@ -442,14 +442,22 @@ class PostgresSaver(BasePostgresSaver):
|
||||
including its configuration, metadata, parent checkpoint (if any),
|
||||
and pending writes.
|
||||
"""
|
||||
with self._cursor() as cur:
|
||||
channel_values = self._load_blobs(
|
||||
value["channel_values"],
|
||||
thread_id=value["thread_id"],
|
||||
checkpoint_ns=value["checkpoint_ns"],
|
||||
checkpoint_id=value["checkpoint_id"],
|
||||
cur=cur,
|
||||
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": {
|
||||
|
||||
@@ -13,8 +13,7 @@ from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
DeltaChannelSentinel,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -393,80 +392,55 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _aload_delta_chain(
|
||||
async def _aget_channel_writes_cur(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
cur: Any,
|
||||
) -> DeltaChainValue:
|
||||
"""Fetch the full delta chain for a channel in one recursive CTE query (async)."""
|
||||
) -> list[Any]:
|
||||
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async)."""
|
||||
await cur.execute(
|
||||
"""
|
||||
WITH RECURSIVE chain AS (
|
||||
SELECT
|
||||
c.checkpoint_id,
|
||||
c.parent_checkpoint_id,
|
||||
cb.version,
|
||||
cb.type,
|
||||
cb.blob
|
||||
FROM checkpoints c
|
||||
JOIN checkpoint_blobs cb
|
||||
ON cb.thread_id = %s
|
||||
AND cb.checkpoint_ns = %s
|
||||
AND cb.channel = %s
|
||||
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
|
||||
WHERE c.thread_id = %s
|
||||
AND c.checkpoint_ns = %s
|
||||
AND c.checkpoint_id = %s
|
||||
|
||||
WITH RECURSIVE chain(cid, depth) AS (
|
||||
SELECT parent_checkpoint_id, 0
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
|
||||
UNION ALL
|
||||
|
||||
SELECT
|
||||
c.checkpoint_id,
|
||||
c.parent_checkpoint_id,
|
||||
cb.version,
|
||||
cb.type,
|
||||
cb.blob
|
||||
FROM chain prev
|
||||
JOIN checkpoints c ON c.checkpoint_id = prev.parent_checkpoint_id
|
||||
JOIN checkpoint_blobs cb
|
||||
ON cb.thread_id = %s
|
||||
AND cb.checkpoint_ns = %s
|
||||
AND cb.channel = %s
|
||||
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
|
||||
WHERE prev.parent_checkpoint_id IS NOT NULL
|
||||
AND prev.type = 'delta'
|
||||
SELECT c.parent_checkpoint_id, ch.depth + 1
|
||||
FROM checkpoints c
|
||||
JOIN chain ch ON c.checkpoint_id = ch.cid
|
||||
WHERE ch.cid IS NOT NULL
|
||||
)
|
||||
SELECT DISTINCT ON (version) type, blob
|
||||
FROM chain
|
||||
ORDER BY version ASC
|
||||
SELECT cw.type, cw.blob
|
||||
FROM checkpoint_writes cw
|
||||
JOIN chain ON cw.checkpoint_id = chain.cid
|
||||
WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s
|
||||
ORDER BY chain.depth DESC, cw.task_id, cw.idx
|
||||
""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
channel,
|
||||
),
|
||||
)
|
||||
rows = await cur.fetchall()
|
||||
base = None
|
||||
deltas: list[list[Any]] = []
|
||||
for row in rows:
|
||||
blob = self.serde.loads_typed((row["type"], row["blob"]))
|
||||
if isinstance(blob, DeltaValue):
|
||||
deltas.append(blob.delta)
|
||||
else:
|
||||
base = blob
|
||||
return DeltaChainValue(base=base, deltas=deltas)
|
||||
return [self.serde.loads_typed((row["type"], row["blob"])) for row in rows]
|
||||
|
||||
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:
|
||||
"""
|
||||
@@ -489,12 +463,14 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
if blob_values:
|
||||
channel_values = self._load_blobs(blob_values)
|
||||
delta_channels = [
|
||||
k.decode() for k, t, _ in blob_values if t.decode() == "delta"
|
||||
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._aload_delta_chain(
|
||||
channel_values[channel] = await self._aget_channel_writes_cur(
|
||||
thread_id, checkpoint_ns, checkpoint_id, channel, cur
|
||||
)
|
||||
|
||||
|
||||
@@ -11,8 +11,7 @@ from langgraph.checkpoint.base import (
|
||||
WRITES_IDX_MAP,
|
||||
BaseCheckpointSaver,
|
||||
ChannelVersions,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
DeltaChannelSentinel,
|
||||
get_checkpoint_id,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
@@ -188,105 +187,72 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
def _load_blobs(
|
||||
self,
|
||||
blob_values: list[tuple[bytes, bytes, bytes]],
|
||||
*,
|
||||
thread_id: str = "",
|
||||
checkpoint_ns: str = "",
|
||||
checkpoint_id: str = "",
|
||||
cur: Any = None,
|
||||
blob_values: Any,
|
||||
) -> dict[str, Any]:
|
||||
if not blob_values:
|
||||
return {}
|
||||
result: dict[str, Any] = {}
|
||||
delta_channels: list[str] = []
|
||||
for k, t, v in blob_values:
|
||||
channel = k.decode()
|
||||
type_tag = t.decode()
|
||||
if type_tag == "delta":
|
||||
delta_channels.append(channel)
|
||||
elif type_tag != "empty":
|
||||
result[channel] = self.serde.loads_typed((type_tag, v))
|
||||
if delta_channels and cur is not None and checkpoint_id:
|
||||
for channel in delta_channels:
|
||||
result[channel] = self._load_delta_chain(
|
||||
thread_id, checkpoint_ns, checkpoint_id, channel, cur
|
||||
)
|
||||
if type_tag != "empty":
|
||||
result[k.decode()] = self.serde.loads_typed((type_tag, v))
|
||||
return result
|
||||
|
||||
def _load_delta_chain(
|
||||
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,
|
||||
) -> DeltaChainValue:
|
||||
"""Fetch the full delta chain for a channel in one recursive CTE query."""
|
||||
) -> list[Any]:
|
||||
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest."""
|
||||
cur.execute(
|
||||
"""
|
||||
WITH RECURSIVE chain AS (
|
||||
SELECT
|
||||
c.checkpoint_id,
|
||||
c.parent_checkpoint_id,
|
||||
cb.version,
|
||||
cb.type,
|
||||
cb.blob
|
||||
FROM checkpoints c
|
||||
JOIN checkpoint_blobs cb
|
||||
ON cb.thread_id = %s
|
||||
AND cb.checkpoint_ns = %s
|
||||
AND cb.channel = %s
|
||||
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
|
||||
WHERE c.thread_id = %s
|
||||
AND c.checkpoint_ns = %s
|
||||
AND c.checkpoint_id = %s
|
||||
|
||||
WITH RECURSIVE chain(cid, depth) AS (
|
||||
SELECT parent_checkpoint_id, 0
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
|
||||
UNION ALL
|
||||
|
||||
SELECT
|
||||
c.checkpoint_id,
|
||||
c.parent_checkpoint_id,
|
||||
cb.version,
|
||||
cb.type,
|
||||
cb.blob
|
||||
FROM chain prev
|
||||
JOIN checkpoints c ON c.checkpoint_id = prev.parent_checkpoint_id
|
||||
JOIN checkpoint_blobs cb
|
||||
ON cb.thread_id = %s
|
||||
AND cb.checkpoint_ns = %s
|
||||
AND cb.channel = %s
|
||||
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
|
||||
WHERE prev.parent_checkpoint_id IS NOT NULL
|
||||
AND prev.type = 'delta'
|
||||
SELECT c.parent_checkpoint_id, ch.depth + 1
|
||||
FROM checkpoints c
|
||||
JOIN chain ch ON c.checkpoint_id = ch.cid
|
||||
WHERE ch.cid IS NOT NULL
|
||||
)
|
||||
SELECT DISTINCT ON (version) type, blob
|
||||
FROM chain
|
||||
ORDER BY version ASC
|
||||
SELECT cw.type, cw.blob
|
||||
FROM checkpoint_writes cw
|
||||
JOIN chain ON cw.checkpoint_id = chain.cid
|
||||
WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s
|
||||
ORDER BY chain.depth DESC, cw.task_id, cw.idx
|
||||
""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
channel,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channel,
|
||||
channel,
|
||||
),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
base = None
|
||||
deltas: list[list[Any]] = []
|
||||
for row in rows:
|
||||
blob = self.serde.loads_typed((row["type"], row["blob"]))
|
||||
if isinstance(blob, DeltaValue):
|
||||
deltas.append(blob.delta)
|
||||
else:
|
||||
base = blob
|
||||
return DeltaChainValue(base=base, deltas=deltas)
|
||||
return [
|
||||
self.serde.loads_typed((row["type"], row["blob"])) for row in cur.fetchall()
|
||||
]
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
|
||||
@@ -3,10 +3,12 @@ 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,
|
||||
@@ -32,19 +34,17 @@ PendingWrite = tuple[str, str, Any]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DeltaValue:
|
||||
"""Returned by DeltaChannel.checkpoint(). Represents one step's writes."""
|
||||
class DeltaChannelSentinel:
|
||||
"""Marker stored in checkpoint_blobs for a DeltaChannel field.
|
||||
|
||||
delta: list[Any]
|
||||
No data is stored here — the actual per-step writes live in checkpoint_writes
|
||||
and are replayed through the reducer at load time.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DeltaChainValue:
|
||||
"""Passed to DeltaChannel.from_checkpoint(). Assembled by the pregel layer."""
|
||||
|
||||
base: list[Any] | None # starting accumulated value; None = start from empty
|
||||
deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest
|
||||
|
||||
_DELTA_RECONSTRUCTION: threading.local = threading.local()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -475,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,8 +20,7 @@ from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
DeltaChannelSentinel,
|
||||
SerializerProtocol,
|
||||
get_checkpoint_id,
|
||||
get_checkpoint_metadata,
|
||||
@@ -127,68 +126,56 @@ class InMemorySaver(
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
versions: ChannelVersions,
|
||||
checkpoint_id: str = "",
|
||||
) -> dict[str, Any]:
|
||||
channel_values: dict[str, Any] = {}
|
||||
delta_channels: list[str] = []
|
||||
for k, v in versions.items():
|
||||
kk = (thread_id, checkpoint_ns, k, v)
|
||||
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] == "delta":
|
||||
delta_channels.append(k)
|
||||
elif vv[0] != "empty":
|
||||
channel_values[k] = self.serde.loads_typed(vv)
|
||||
for channel in delta_channels:
|
||||
channel_values[channel] = self._assemble_delta_chain(
|
||||
thread_id, checkpoint_ns, checkpoint_id, channel
|
||||
)
|
||||
return channel_values
|
||||
if vv[0] == "empty":
|
||||
continue
|
||||
result[k] = self.serde.loads_typed(vv)
|
||||
return result
|
||||
|
||||
def _assemble_delta_chain(
|
||||
def _resolve_delta_channels(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> DeltaChainValue:
|
||||
"""Walk the checkpoint parent tree to collect all delta blobs for a channel."""
|
||||
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, {})
|
||||
blobs: list[Any] = []
|
||||
current_id: str | None = checkpoint_id
|
||||
seen_versions: set[str] = set()
|
||||
while current_id is not None:
|
||||
entry = ns_storage.get(current_id)
|
||||
# Walk the parent chain, collecting checkpoint IDs oldest→newest.
|
||||
chain: list[str] = []
|
||||
current: str | None = checkpoint_id
|
||||
while current is not None:
|
||||
entry = ns_storage.get(current)
|
||||
if entry is None:
|
||||
break
|
||||
checkpoint = self.serde.loads_typed(entry[0])
|
||||
version = checkpoint["channel_versions"].get(channel)
|
||||
if version is None:
|
||||
break # channel not yet in this checkpoint
|
||||
_, _, current_id = entry # advance to parent before the continue/break
|
||||
if version in seen_versions:
|
||||
continue # same blob already collected; keep walking to older checkpoints
|
||||
seen_versions.add(version)
|
||||
kk = (thread_id, checkpoint_ns, channel, version)
|
||||
if kk not in self.blobs:
|
||||
break
|
||||
vv = self.blobs[kk]
|
||||
if vv[0] == "empty":
|
||||
break
|
||||
blob = self.serde.loads_typed(vv)
|
||||
blobs.append(blob)
|
||||
if not isinstance(blob, DeltaValue):
|
||||
break # hit a snapshot (plain list) — chain root found
|
||||
blobs.reverse()
|
||||
base = None
|
||||
deltas: list[list[Any]] = []
|
||||
for blob in blobs:
|
||||
if isinstance(blob, DeltaValue):
|
||||
deltas.append(blob.delta)
|
||||
else:
|
||||
base = blob
|
||||
return DeltaChainValue(base=base, deltas=deltas)
|
||||
chain.append(current)
|
||||
_, _, parent = entry
|
||||
current = parent
|
||||
# Collect writes for `channel` from each checkpoint in oldest→newest order.
|
||||
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.
|
||||
@@ -211,16 +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"],
|
||||
checkpoint_id,
|
||||
),
|
||||
"channel_values": channel_values,
|
||||
},
|
||||
metadata=self.serde.loads_typed(metadata),
|
||||
pending_writes=[
|
||||
@@ -244,22 +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"],
|
||||
checkpoint_id,
|
||||
),
|
||||
"channel_values": channel_values,
|
||||
},
|
||||
metadata=self.serde.loads_typed(metadata),
|
||||
pending_writes=[
|
||||
@@ -354,22 +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"],
|
||||
checkpoint_id,
|
||||
),
|
||||
"channel_values": channel_values,
|
||||
},
|
||||
metadata=metadata,
|
||||
parent_config=(
|
||||
|
||||
@@ -80,8 +80,6 @@ SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
|
||||
("langgraph.types", "Overwrite"),
|
||||
("langgraph.store.base", "Item"),
|
||||
("langgraph.store.base", "GetOp"),
|
||||
# DeltaChannel checkpoint value type
|
||||
("langgraph.checkpoint.base", "DeltaValue"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -64,10 +64,12 @@ def _warn_once(
|
||||
logger.warning(msg, *args)
|
||||
|
||||
|
||||
def _is_delta_value(obj: Any) -> bool:
|
||||
from langgraph.checkpoint.base import DeltaValue # lazy import avoids circular dep
|
||||
def _get_delta_sentinel_cls() -> type:
|
||||
from langgraph.checkpoint.base import (
|
||||
DeltaChannelSentinel,
|
||||
) # lazy import avoids circular dep
|
||||
|
||||
return isinstance(obj, DeltaValue)
|
||||
return DeltaChannelSentinel
|
||||
|
||||
|
||||
class JsonPlusSerializer(SerializerProtocol):
|
||||
@@ -262,8 +264,8 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
return "bytes", obj
|
||||
elif isinstance(obj, bytearray):
|
||||
return "bytearray", obj
|
||||
elif _is_delta_value(obj):
|
||||
return "delta", _msgpack_enc({"d": obj.delta})
|
||||
elif isinstance(obj, _get_delta_sentinel_cls()):
|
||||
return "delta", b""
|
||||
else:
|
||||
try:
|
||||
return "msgpack", _msgpack_enc(obj)
|
||||
@@ -287,12 +289,9 @@ class JsonPlusSerializer(SerializerProtocol):
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
elif type_ == "delta":
|
||||
from langgraph.checkpoint.base import DeltaValue # lazy import
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
|
||||
raw = ormsgpack.unpackb(
|
||||
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
return DeltaValue(delta=raw["d"])
|
||||
return DeltaChannelSentinel()
|
||||
elif self.pickle_fallback and type_ == "pickle":
|
||||
return pickle.loads(data_)
|
||||
else:
|
||||
|
||||
@@ -999,14 +999,13 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_delta_value_serde_round_trip() -> None:
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
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 = DeltaValue(delta=[{"type": "human", "content": "hi"}])
|
||||
original = DeltaChannelSentinel()
|
||||
type_tag, blob = serde.dumps_typed(original)
|
||||
assert type_tag == "delta"
|
||||
loaded = serde.loads_typed((type_tag, blob))
|
||||
assert isinstance(loaded, DeltaValue)
|
||||
assert loaded.delta == original.delta
|
||||
assert isinstance(loaded, DeltaChannelSentinel)
|
||||
|
||||
@@ -323,11 +323,10 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
|
||||
|
||||
|
||||
class TestInMemorySaverDeltaChannel:
|
||||
def test_load_blobs_assembles_delta_chain(self) -> None:
|
||||
"""_load_blobs returns DeltaChainValue for delta channels, not raw DeltaValue."""
|
||||
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 (
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
DeltaChannelSentinel,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
@@ -336,55 +335,59 @@ class TestInMemorySaverDeltaChannel:
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
v1 = "00000000000000000000000000000001.0000000000000000"
|
||||
v2 = "00000000000000000000000000000002.0000000000000000"
|
||||
|
||||
delta1 = DeltaValue(delta=[{"content": "hi"}])
|
||||
delta2 = DeltaValue(delta=[{"content": "bye"}])
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(delta1)
|
||||
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(delta2)
|
||||
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"
|
||||
cp2["channel_versions"][channel] = v2
|
||||
saver.storage[thread_id][ns] = {
|
||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
||||
}
|
||||
|
||||
result = saver._load_blobs(thread_id, ns, {channel: v2}, "cp2")
|
||||
assert channel in result
|
||||
chain = result[channel]
|
||||
assert isinstance(chain, DeltaChainValue)
|
||||
assert chain.deltas == [[{"content": "hi"}], [{"content": "bye"}]]
|
||||
|
||||
def test_load_blobs_single_delta_no_parent(self) -> None:
|
||||
"""Single delta with no parent checkpoint produces a chain with one delta."""
|
||||
from langgraph.checkpoint.base import (
|
||||
DeltaChainValue,
|
||||
DeltaValue,
|
||||
empty_checkpoint,
|
||||
# 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"}),
|
||||
"",
|
||||
)
|
||||
|
||||
saver = InMemorySaver()
|
||||
serde = JsonPlusSerializer()
|
||||
|
||||
thread_id, ns, channel = "t1", "", "messages"
|
||||
v1 = "00000000000000000000000000000001.0000000000000000"
|
||||
delta = DeltaValue(delta=[{"content": "only"}])
|
||||
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(delta)
|
||||
|
||||
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)
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": ns,
|
||||
"checkpoint_id": "cp2",
|
||||
}
|
||||
}
|
||||
|
||||
result = saver._load_blobs(thread_id, ns, {channel: v1}, "cp1")
|
||||
chain = result[channel]
|
||||
assert isinstance(chain, DeltaChainValue)
|
||||
assert chain.base is None
|
||||
assert chain.deltas == [[{"content": "only"}]]
|
||||
result = saver.get_channel_writes(config, channel)
|
||||
assert result == [{"content": "hi"}, {"content": "bye"}]
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
from langgraph.checkpoint.base import DeltaChannelSentinel
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
@@ -14,64 +14,43 @@ from langgraph.errors import EmptyChannelError
|
||||
__all__ = ("DeltaChannel",)
|
||||
|
||||
|
||||
class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
|
||||
"""A channel that stores only per-step write deltas in checkpoints.
|
||||
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.
|
||||
|
||||
Reconstructs the full accumulated list at load time by replaying the
|
||||
chain of deltas through the operator. Use with append-style reducers
|
||||
(e.g. `add_messages`) on long-running threads to reduce checkpoint
|
||||
storage from O(N²) to O(N).
|
||||
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 a dedicated blob store
|
||||
(InMemorySaver, PostgresSaver) use an O(1) fast-path per chain step;
|
||||
all others (SQLite, MongoDB, etc.) fall back to get_tuple traversal.
|
||||
|
||||
Use `snapshot_every=N` to cap chain traversal depth at N steps. Every N
|
||||
steps a full snapshot is written as the chain root; subsequent deltas
|
||||
chain back to it, so `get_state` / reload never traverses more than N
|
||||
checkpoints regardless of thread length. Recommended for savers without
|
||||
a dedicated blob store.
|
||||
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)]
|
||||
# Cap reconstruction depth (recommended for SQLite / MongoDB savers):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages, snapshot_every=50)]
|
||||
# Dict-type reducer (type inferred from the Annotated outer type):
|
||||
files: Annotated[dict, DeltaChannel(merge_files)]
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"value",
|
||||
"operator",
|
||||
"snapshot_every",
|
||||
"_pending",
|
||||
"_base_version",
|
||||
"_overwritten",
|
||||
"_steps_since_snapshot",
|
||||
)
|
||||
__slots__ = ("value", "operator")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
operator: Callable[[list[Value], Any], list[Value]],
|
||||
*,
|
||||
snapshot_every: int | None = None,
|
||||
snapshot_every: int | None = None, # reserved for future use
|
||||
) -> None:
|
||||
super().__init__(list)
|
||||
self.operator = operator
|
||||
self.snapshot_every = snapshot_every
|
||||
self.value: list[Value] = []
|
||||
self._pending: list[Any] = []
|
||||
self._base_version: str | None = None
|
||||
self._overwritten: bool = False
|
||||
self._steps_since_snapshot: int = 0
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DeltaChannel):
|
||||
return False
|
||||
if self.snapshot_every != other.snapshot_every:
|
||||
return False
|
||||
if (
|
||||
self.operator.__name__ != "<lambda>"
|
||||
and other.operator.__name__ != "<lambda>"
|
||||
@@ -88,18 +67,14 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
|
||||
return self.typ | list[self.typ] # type: ignore[name-defined]
|
||||
|
||||
def copy(self) -> Self:
|
||||
new = DeltaChannel(self.operator, snapshot_every=self.snapshot_every)
|
||||
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()
|
||||
new._pending = self._pending[:]
|
||||
new._base_version = self._base_version
|
||||
new._overwritten = self._overwritten
|
||||
new._steps_since_snapshot = self._steps_since_snapshot
|
||||
return new
|
||||
|
||||
def from_checkpoint(self, checkpoint: Any) -> Self:
|
||||
new = DeltaChannel(self.operator, snapshot_every=self.snapshot_every)
|
||||
new = DeltaChannel(self.operator)
|
||||
new.typ = self.typ
|
||||
new.key = self.key
|
||||
if checkpoint is MISSING:
|
||||
@@ -107,29 +82,18 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
|
||||
new.value = new.typ()
|
||||
except Exception:
|
||||
new.value = []
|
||||
elif isinstance(checkpoint, DeltaChainValue):
|
||||
accumulated: list[Value] = (
|
||||
checkpoint.base if checkpoint.base is not None else new.typ()
|
||||
)
|
||||
for step_writes in checkpoint.deltas:
|
||||
for write in step_writes:
|
||||
accumulated = new.operator(accumulated, write)
|
||||
new.value = accumulated
|
||||
# Seed the counter from actual chain depth so rehydration fires at
|
||||
# the right time regardless of how many prior invocations there were.
|
||||
new._steps_since_snapshot = len(checkpoint.deltas)
|
||||
elif isinstance(checkpoint, DeltaValue):
|
||||
raise ValueError(
|
||||
f"Channel '{self.key}' uses DeltaChannel but the checkpointer "
|
||||
"does not support incremental channel storage. "
|
||||
"Use InMemorySaver or PostgresSaver, or remove DeltaChannel from your schema."
|
||||
)
|
||||
elif isinstance(checkpoint, list):
|
||||
# Flat list of individual write values (oldest→newest) from get_channel_writes.
|
||||
value: Any = new.typ()
|
||||
for write in checkpoint:
|
||||
value = new.operator(value, write)
|
||||
new.value = value
|
||||
else:
|
||||
# Backwards compat: plain list from old BinaryOperatorAggregate checkpoint.
|
||||
new.value = list(checkpoint)
|
||||
new._pending = []
|
||||
new._base_version = None # set by the subsequent after_checkpoint() call
|
||||
new._overwritten = False
|
||||
# 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:
|
||||
@@ -154,13 +118,10 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
|
||||
self.value = (
|
||||
list(overwrite_value) if overwrite_value is not None else self.typ()
|
||||
)
|
||||
self._pending = list(self.value)
|
||||
self._overwritten = True
|
||||
seen_overwrite = True
|
||||
elif not seen_overwrite:
|
||||
base = self.typ() if self.value is MISSING else self.value
|
||||
self.value = self.operator(base, value)
|
||||
self._pending.append(value)
|
||||
return True
|
||||
|
||||
def get(self) -> list[Value]:
|
||||
@@ -171,26 +132,5 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
|
||||
def is_available(self) -> bool:
|
||||
return self.value is not MISSING
|
||||
|
||||
def checkpoint(self) -> Any:
|
||||
if (
|
||||
self.snapshot_every is not None
|
||||
and self._steps_since_snapshot >= self.snapshot_every
|
||||
):
|
||||
# Emit a full snapshot to cap chain depth at snapshot_every.
|
||||
# The saver stores this as a plain (non-diff) blob, so future
|
||||
# deltas will chain back to it and traversal depth resets to 1.
|
||||
return list(self.value)
|
||||
return DeltaValue(delta=self._pending[:])
|
||||
|
||||
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
|
||||
if version != self._base_version:
|
||||
if self._base_version is None:
|
||||
pass # First call after from_checkpoint — anchor without counting a step.
|
||||
elif self.snapshot_every is not None:
|
||||
if self._steps_since_snapshot >= self.snapshot_every:
|
||||
self._steps_since_snapshot = 0
|
||||
else:
|
||||
self._steps_since_snapshot += 1
|
||||
self._base_version = version
|
||||
self._pending = []
|
||||
self._overwritten = False
|
||||
def checkpoint(self) -> DeltaChannelSentinel:
|
||||
return DeltaChannelSentinel()
|
||||
|
||||
@@ -121,26 +121,22 @@ def test_untracked_value() -> None:
|
||||
|
||||
def test_delta_channel_basic_two_steps() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
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)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
# Step 1: one message added
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
assert len(d1.delta) == 1
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
assert isinstance(d1, DeltaChannelSentinel)
|
||||
|
||||
# Step 2: another message
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert len(d2.delta) == 1
|
||||
ch.after_checkpoint("v2")
|
||||
assert isinstance(d2, DeltaChannelSentinel)
|
||||
|
||||
# Full accumulated value is preserved in memory
|
||||
assert len(ch.get()) == 2
|
||||
@@ -148,40 +144,21 @@ def test_delta_channel_basic_two_steps() -> None:
|
||||
assert ch.get()[1].content == "hello"
|
||||
|
||||
|
||||
def test_delta_channel_after_checkpoint_no_op_when_unchanged() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
ch.after_checkpoint("v1")
|
||||
|
||||
# Same version: no-op
|
||||
ch.after_checkpoint("v1")
|
||||
assert ch._base_version == "v1"
|
||||
assert ch._pending == []
|
||||
|
||||
|
||||
def test_delta_channel_from_checkpoint_chain() -> None:
|
||||
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.checkpoint.base import DeltaChainValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
chain = DeltaChainValue(
|
||||
base=None,
|
||||
deltas=[
|
||||
[HumanMessage(content="hi", id="h1")],
|
||||
[AIMessage(content="hello", id="a1")],
|
||||
[HumanMessage(content="bye", id="h2")],
|
||||
],
|
||||
)
|
||||
ch = spec.from_checkpoint(chain)
|
||||
# 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"
|
||||
@@ -195,68 +172,45 @@ def test_delta_channel_from_checkpoint_backwards_compat() -> None:
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# Old BinaryOperatorAggregate checkpoint: plain list
|
||||
# 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_resets_chain() -> None:
|
||||
def test_delta_channel_overwrite() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
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.after_checkpoint(None)
|
||||
ch.update([HumanMessage(content="old", id="h1")])
|
||||
ch.after_checkpoint("v1")
|
||||
|
||||
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
|
||||
d = ch.checkpoint()
|
||||
assert isinstance(d, DeltaValue)
|
||||
assert len(d.delta) == 1
|
||||
assert d.delta[0].content == "new"
|
||||
# _overwritten flag must be set so next checkpoint acts as a chain root
|
||||
assert ch._overwritten is True
|
||||
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_unsupported_saver_raises() -> None:
|
||||
"""from_checkpoint raises ValueError when the saver returns a raw DeltaValue."""
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
raw = DeltaValue(delta=[{"type": "human", "content": "hello"}])
|
||||
with pytest.raises(
|
||||
ValueError, match="does not support incremental channel storage"
|
||||
):
|
||||
spec.from_checkpoint(raw)
|
||||
|
||||
|
||||
def test_delta_channel_remove_message_delta_and_replay() -> None:
|
||||
"""RemoveMessage stored in a delta must round-trip correctly through the chain."""
|
||||
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.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
# Step 1: add two messages
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
ch.update([AIMessage(content="hello", id="a1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
assert ch.get() == [
|
||||
HumanMessage(content="hi", id="h1"),
|
||||
AIMessage(content="hello", id="a1"),
|
||||
@@ -264,127 +218,63 @@ def test_delta_channel_remove_message_delta_and_replay() -> None:
|
||||
|
||||
# Step 2: remove the AI message
|
||||
ch.update([RemoveMessage(id="a1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert isinstance(d2, DeltaValue)
|
||||
assert any(isinstance(w, RemoveMessage) for w in d2.delta)
|
||||
ch.after_checkpoint("v2", checkpoint_id="cid2")
|
||||
assert ch.get() == [HumanMessage(content="hi", id="h1")]
|
||||
|
||||
# Replay the full chain from scratch — must reproduce the post-remove state
|
||||
chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta])
|
||||
ch2 = spec.from_checkpoint(chain)
|
||||
# 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_delta_and_replay() -> None:
|
||||
"""Updating a message by ID stored in a delta must round-trip correctly."""
|
||||
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.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
# Step 1: add a message
|
||||
ch.update([HumanMessage(content="original", id="h1")])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
# Step 2: update the same message by ID
|
||||
ch.update([HumanMessage(content="updated", id="h1")])
|
||||
d2 = ch.checkpoint()
|
||||
assert isinstance(d2, DeltaValue)
|
||||
ch.after_checkpoint("v2", checkpoint_id="cid2")
|
||||
assert ch.get() == [HumanMessage(content="updated", id="h1")]
|
||||
|
||||
# Replay the full chain — must produce the updated message, not the original
|
||||
chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta])
|
||||
ch2 = spec.from_checkpoint(chain)
|
||||
# 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_snapshot_every_emits_plain_list() -> None:
|
||||
"""snapshot_every=N causes a plain-list snapshot after N steps; next deltas chain to it."""
|
||||
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
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
SNAP = 3
|
||||
spec = DeltaChannel(add_messages, snapshot_every=SNAP)
|
||||
ch = spec.from_checkpoint(MISSING)
|
||||
# First after_checkpoint anchors _base_version without counting a step.
|
||||
ch.after_checkpoint("v0", checkpoint_id="cid0")
|
||||
|
||||
# Steps 1..SNAP: each should stay as DeltaValue; counter increments each step.
|
||||
for i in range(1, SNAP + 1):
|
||||
ch.update([HumanMessage(content=f"m{i}", id=f"h{i}")])
|
||||
ckpt = ch.checkpoint()
|
||||
assert isinstance(ckpt, DeltaValue), f"expected DeltaValue at step {i}"
|
||||
ch.after_checkpoint(f"v{i}", checkpoint_id=f"cid{i}")
|
||||
|
||||
# Step SNAP+1: _steps_since_snapshot == SNAP → snapshot fires
|
||||
ch.update([HumanMessage(content="snap", id="hsnap")])
|
||||
snap = ch.checkpoint()
|
||||
assert isinstance(snap, list), "expected plain-list snapshot at snapshot_every step"
|
||||
assert len(snap) == SNAP + 1
|
||||
|
||||
# After snapshot, counter resets — next step is DeltaValue again
|
||||
ch.after_checkpoint("vsnap", checkpoint_id="cidsnap")
|
||||
ch.update([HumanMessage(content="post", id="hpost")])
|
||||
post = ch.checkpoint()
|
||||
assert isinstance(post, DeltaValue)
|
||||
ch.update([HumanMessage(content="hi", id="h1")])
|
||||
assert isinstance(ch.checkpoint(), DeltaChannelSentinel)
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_every_end_to_end() -> None:
|
||||
"""Graph with snapshot_every: get_state returns correct accumulated value after snapshot."""
|
||||
from typing import Annotated
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)]
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
def respond(state: State) -> dict:
|
||||
counter["n"] += 1
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
|
||||
]
|
||||
}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("respond", respond)
|
||||
builder.add_edge(START, "respond")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "snap-test"}}
|
||||
|
||||
# Run 5 turns — snapshot fires after 2 steps, then again after 2 more
|
||||
for i in range(5):
|
||||
graph.invoke({"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config)
|
||||
|
||||
state = graph.get_state(config)
|
||||
msgs = state.values["messages"]
|
||||
# 5 human + 5 AI = 10 total
|
||||
assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}"
|
||||
|
||||
|
||||
def test_delta_channel_inmemory_saver_assembles_chain() -> None:
|
||||
"""InMemorySaver assembles the delta chain inside get_tuple (no pregel involvement)."""
|
||||
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
|
||||
@@ -414,14 +304,16 @@ def test_delta_channel_inmemory_saver_assembles_chain() -> None:
|
||||
graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
|
||||
graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config)
|
||||
|
||||
# get_tuple must return a fully assembled DeltaChainValue, not raw DeltaValue
|
||||
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
|
||||
# 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"], DeltaValue)
|
||||
assert isinstance(saved.checkpoint["channel_values"]["messages"], DeltaChainValue)
|
||||
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
|
||||
@@ -451,48 +343,39 @@ def test_delta_channel_dict_reducer_fresh_channel() -> None:
|
||||
|
||||
def test_delta_channel_dict_reducer_basic_updates() -> None:
|
||||
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
|
||||
from langgraph.checkpoint.base import DeltaValue
|
||||
from langgraph.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.after_checkpoint(None)
|
||||
|
||||
ch.update([{"a": 1}])
|
||||
d1 = ch.checkpoint()
|
||||
assert isinstance(d1, DeltaValue)
|
||||
assert d1.delta == [{"a": 1}]
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
assert isinstance(d1, DeltaChannelSentinel)
|
||||
|
||||
ch.update([{"b": 2}])
|
||||
d2 = ch.checkpoint()
|
||||
assert d2.delta == [{"b": 2}]
|
||||
ch.after_checkpoint("v2")
|
||||
assert isinstance(d2, DeltaChannelSentinel)
|
||||
|
||||
assert ch.get() == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_chain_reconstruction() -> None:
|
||||
"""DeltaChainValue replays correctly through a dict merge reducer."""
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
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)
|
||||
chain = DeltaChainValue(
|
||||
base={"a": 1},
|
||||
deltas=[[{"b": 2}], [{"c": 3}]],
|
||||
)
|
||||
ch = spec.from_checkpoint(chain)
|
||||
# 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}
|
||||
assert ch._steps_since_snapshot == 2
|
||||
|
||||
|
||||
def test_delta_channel_dict_reducer_with_deletions() -> None:
|
||||
"""Dict reducer that treats None values as deletions works end-to-end (deepagents pattern)."""
|
||||
from langgraph.checkpoint.base import DeltaChainValue
|
||||
|
||||
def merge_files(left: dict | None, right: dict) -> dict:
|
||||
if left is None:
|
||||
@@ -506,25 +389,19 @@ def test_delta_channel_dict_reducer_with_deletions() -> None:
|
||||
return result
|
||||
|
||||
ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING)
|
||||
ch.after_checkpoint(None)
|
||||
|
||||
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
|
||||
ch.after_checkpoint("v1", checkpoint_id="cid1")
|
||||
|
||||
# Delete file1, add file3
|
||||
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
||||
ch.after_checkpoint("v2", checkpoint_id="cid2")
|
||||
|
||||
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
# Confirm chain reconstruction produces the same result
|
||||
chain = DeltaChainValue(
|
||||
base={},
|
||||
deltas=[
|
||||
[{"file1.py": "content1", "file2.py": "content2"}],
|
||||
[{"file1.py": None, "file3.py": "content3"}],
|
||||
],
|
||||
)
|
||||
# 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(chain)
|
||||
ch2 = spec.from_checkpoint(writes)
|
||||
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
@@ -8,6 +8,11 @@ Simulates realistic multi-turn conversations with paragraph-length messages
|
||||
|
||||
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
|
||||
@@ -42,8 +47,6 @@ try:
|
||||
except ImportError:
|
||||
_POSTGRES_AVAILABLE = False
|
||||
|
||||
SNAPSHOT_EVERY = 50
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Realistic message payload (~100 tokens / ~400 chars each)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -121,10 +124,6 @@ class DeltaState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages)]
|
||||
|
||||
|
||||
class DeltaSnapshotState(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=SNAPSHOT_EVERY)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graph factory
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -239,7 +238,12 @@ def run_benchmark() -> None:
|
||||
|
||||
checkpointers: list[tuple[str, Any]] = [("InMemory", None)]
|
||||
if _POSTGRES_AVAILABLE:
|
||||
checkpointers.append(("Postgres (recursive CTE)", "postgres"))
|
||||
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} ---")
|
||||
@@ -277,30 +281,28 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
|
||||
with _make_saver() as saver:
|
||||
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
|
||||
with _make_saver() as saver:
|
||||
s_wt, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
|
||||
rows.append((turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt))
|
||||
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt))
|
||||
|
||||
# ── Table 1: Storage ─────────────────────────────────────────────────────
|
||||
W = 80
|
||||
W = 70
|
||||
print("Storage (checkpoint blob bytes)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'delta+snap':>12} {'savings':>8}"
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'savings':>8}"
|
||||
)
|
||||
print("-" * W)
|
||||
storage_results = []
|
||||
for turns, b_bytes, d_bytes, s_bytes, *_ in rows:
|
||||
for turns, b_bytes, d_bytes, *_ in rows:
|
||||
if b_bytes < 0:
|
||||
print(
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>12} {'n/a':>8}"
|
||||
f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>8}"
|
||||
)
|
||||
else:
|
||||
ratio = b_bytes / s_bytes if s_bytes else float("inf")
|
||||
storage_results.append((turns, b_bytes, s_bytes, ratio))
|
||||
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} {_fmt_bytes(s_bytes):>12} "
|
||||
f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} "
|
||||
f"{ratio:>7.0f}x"
|
||||
)
|
||||
print("=" * W)
|
||||
@@ -310,35 +312,30 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
|
||||
print("Read latency (avg of 5 get_state calls)")
|
||||
print("=" * W)
|
||||
print(
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'delta+snap':>12}"
|
||||
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}"
|
||||
)
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt in rows:
|
||||
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 {s_rt * 1000:>10.1f}ms"
|
||||
f"{b_rt * 1000:>10.1f}ms {d_rt * 1000:>10.1f}ms"
|
||||
)
|
||||
print("=" * W)
|
||||
print()
|
||||
|
||||
if storage_results:
|
||||
best = storage_results[-1]
|
||||
turns, b_bytes, s_bytes, ratio = best
|
||||
_, _, _, _, b_rt, _, s_rt = rows[-1]
|
||||
turns, b_bytes, d_bytes, ratio = best
|
||||
_, _, _, b_rt, d_rt = rows[-1]
|
||||
print(
|
||||
f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(s_bytes)} ({ratio:.0f}x less storage); "
|
||||
f"read {b_rt * 1000:.1f}ms → {s_rt * 1000:.1f}ms"
|
||||
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, unbounded chain"
|
||||
)
|
||||
print(
|
||||
f" delta+snap = DeltaChannel(add_messages, snapshot_every={SNAPSHOT_EVERY}) — O(N) storage, bounded read depth"
|
||||
)
|
||||
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
|
||||
print(" delta = DeltaChannel(add_messages) — O(N) storage, reconstructed from writes")
|
||||
print()
|
||||
|
||||
|
||||
@@ -359,15 +356,10 @@ def test_delta_channel_benchmark(capsys: Any) -> None:
|
||||
for turns in [25, 50]:
|
||||
_, _, b_bytes = _run_turns(turns, BinaryState)
|
||||
_, _, d_bytes = _run_turns(turns, DeltaState)
|
||||
_, _, s_bytes = _run_turns(turns, DeltaSnapshotState)
|
||||
assert d_bytes < b_bytes, (
|
||||
f"DeltaChannel should use less storage at {turns} turns, "
|
||||
f"got delta={d_bytes} binary={b_bytes}"
|
||||
)
|
||||
assert s_bytes < b_bytes, (
|
||||
f"DeltaChannel+snapshot should use less storage at {turns} turns, "
|
||||
f"got snapshot={s_bytes} binary={b_bytes}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user