diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 537acd0fa..6e884e578 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -8,6 +8,7 @@ from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, @@ -446,8 +447,6 @@ class PostgresSaver(BasePostgresSaver): including its configuration, metadata, parent checkpoint (if any), and pending writes. """ - from langgraph.checkpoint.base import DELTA_SENTINEL - channel_values = self._load_blobs(value["channel_values"]) if any(v is DELTA_SENTINEL for v in channel_values.values()): cp_config = cast( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 10e3fadcc..de35d0a14 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -402,12 +402,7 @@ class AsyncPostgresSaver(BasePostgresSaver): cur: Any, ) -> list[Any]: """Async version of _get_channel_writes_cur — see sync version for rationale.""" - try: - from langgraph.types import ( - Overwrite, # type: ignore[import-untyped,import-not-found] - ) - except ImportError: - Overwrite = None # type: ignore[assignment] + from langgraph.types import Overwrite # type: ignore[import-untyped] await cur.execute( "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " @@ -440,7 +435,7 @@ class AsyncPostgresSaver(BasePostgresSaver): for type_tag, blob in writes_by_cp.get(cid, []): val = self.serde.loads_typed((type_tag, blob)) collected.append(val) - if Overwrite is not None and isinstance(val, Overwrite): + if isinstance(val, Overwrite): collected.reverse() return collected collected.reverse() diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index e4564ed6c..976b36194 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -235,13 +235,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): 1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread. 2. Walk ancestry in Python, then fetch writes with a plain ANY() filter. """ - # Lazy import: mirrors the Send pattern in the serializer. - try: - from langgraph.types import ( - Overwrite, # type: ignore[import-untyped,import-not-found] - ) - except ImportError: - Overwrite = None # type: ignore[assignment] + from langgraph.types import Overwrite # type: ignore[import-untyped] cur.execute( "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " @@ -274,7 +268,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): for type_tag, blob in writes_by_cp.get(cid, []): val = self.serde.loads_typed((type_tag, blob)) collected.append(val) - if Overwrite is not None and isinstance(val, Overwrite): + if isinstance(val, Overwrite): collected.reverse() return collected collected.reverse() diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 9fcacea36..10382519e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -4,9 +4,10 @@ import copy import logging import threading from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence -from typing import ( +from typing import ( # noqa: UP035 Any, Generic, + List, Literal, NamedTuple, TypedDict, @@ -20,13 +21,17 @@ from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_ from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( - DELTA_SENTINEL, + DELTA_SENTINEL as DELTA_SENTINEL, +) +from langgraph.checkpoint.serde.types import ( ERROR, INTERRUPT, RESUME, SCHEDULED, ChannelProtocol, - DeltaChannelWrites, +) +from langgraph.checkpoint.serde.types import ( + DeltaChannelWrites as DeltaChannelWrites, ) V = TypeVar("V", int, float, str) @@ -34,6 +39,7 @@ PendingWrite = tuple[str, str, Any] _DELTA_RECONSTRUCTION: threading.local = threading.local() + logger = logging.getLogger(__name__) @@ -463,7 +469,7 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]: + def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006 """Collect writes for `channel` across this checkpoint's ancestry, oldest→newest. Scans newest→oldest and stops at the first `Overwrite` (either from @@ -471,18 +477,18 @@ class BaseCheckpointSaver(Generic[V]): Default implementation walks the full thread history via `list()`; savers can override with a more efficient query (InMemorySaver and PostgresSaver do this). - """ - try: - from langgraph.types import Overwrite # type: ignore[import-not-found] - except ImportError: - Overwrite = None # type: ignore[assignment] + `List` is used instead of `list` to avoid mypy confusing it with the + saver's own `list` method. + """ # Guard against re-entrant calls: when list() triggers reconstruction # which calls list() again, the inner call returns tuples with # DELTA_SENTINEL 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 [] + from langgraph.types import Overwrite # type: ignore[import-untyped] + _DELTA_RECONSTRUCTION.active = True try: collected: list[Any] = [] # newest first @@ -498,7 +504,7 @@ class BaseCheckpointSaver(Generic[V]): if ch != channel: continue collected.append(value) - if Overwrite is not None and isinstance(value, Overwrite): + if isinstance(value, Overwrite): collected.reverse() return collected collected.reverse() @@ -508,12 +514,9 @@ class BaseCheckpointSaver(Generic[V]): async def aget_channel_writes( self, config: RunnableConfig, channel: str - ) -> list[Any]: + ) -> List[Any]: # noqa: UP006 """Async version of get_channel_writes.""" - try: - from langgraph.types import Overwrite # type: ignore[import-not-found] - except ImportError: - Overwrite = None # type: ignore[assignment] + from langgraph.types import Overwrite # type: ignore[import-untyped] if getattr(_DELTA_RECONSTRUCTION, "active", False): return [] @@ -530,7 +533,7 @@ class BaseCheckpointSaver(Generic[V]): if ch != channel: continue collected.append(value) - if Overwrite is not None and isinstance(value, Overwrite): + if isinstance(value, Overwrite): collected.reverse() return collected collected.reverse() diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index ee992d324..9947be125 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -14,13 +14,13 @@ from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, - DELTA_SENTINEL, DeltaChannelWrites, SerializerProtocol, get_checkpoint_id, @@ -154,13 +154,6 @@ class InMemorySaver( ) def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]: - # Lazy import: avoids a hard dep on `langgraph` at module load time - # (mirrors the Send import pattern in the serializer). - try: - from langgraph.types import Overwrite # type: ignore[import-not-found] - except ImportError: - Overwrite = None # type: ignore[assignment] - thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = config["configurable"].get("checkpoint_id", "") @@ -175,6 +168,8 @@ class InMemorySaver( chain.append(current) _, _, parent = entry current = parent + from langgraph.types import Overwrite # type: ignore[import-untyped] + # Scan writes newest→oldest. Stop at the first `Overwrite` — it # dominates all older history. Either from `snapshot_every` or from # user code: the bound applies the same way. @@ -190,7 +185,7 @@ class InMemorySaver( continue val = self.serde.loads_typed(serialized) collected.append(val) - if Overwrite is not None and isinstance(val, Overwrite): + if isinstance(val, Overwrite): collected.reverse() return collected collected.reverse() diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 7a8ba7484..c8df72960 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -43,6 +43,7 @@ class DeltaChannelWrites: writes: list[Any] + Value = TypeVar("Value", covariant=True) Update = TypeVar("Update", contravariant=True) C = TypeVar("C") diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index 25dbf1d29..0247714eb 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -126,6 +126,15 @@ class DeltaState(TypedDict): messages: Annotated[list, DeltaChannel(add_messages)] +_SNAPSHOT_EVERY = 25 + + +class DeltaSnapshotState(TypedDict): + messages: Annotated[ + list, DeltaChannel(add_messages, snapshot_every=_SNAPSHOT_EVERY) + ] + + # --------------------------------------------------------------------------- # Graph factory # --------------------------------------------------------------------------- @@ -223,6 +232,10 @@ def _approx_tokens(n_turns: int) -> str: # Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window). TURN_COUNTS = [10, 25, 50, 100, 500] +# Snapshot-only counts: pure delta replay becomes painful past 500 turns and +# add_messages storage balloons past 1 GB, so only the snapshot variant runs here. +SNAPSHOT_ONLY_TURN_COUNTS = [1000] + def _checkpointer_factories() -> list[tuple[str, Any]]: """Return (label, context_manager_or_none) pairs for available checkpointers.""" @@ -278,52 +291,78 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: with SqliteSaver.from_conn_string(f.name) as saver: yield saver - rows = [] + rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = [] for turns in TURN_COUNTS: with _make_saver() as saver: - b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver) + _, 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)) + _, d_rt, d_bytes = _run_turns(turns, DeltaState, saver) + with _make_saver() as saver: + _, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver) + rows.append((turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt)) + for turns in SNAPSHOT_ONLY_TURN_COUNTS: + with _make_saver() as saver: + _, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver) + rows.append((turns, None, None, s_bytes, None, None, s_rt)) # ── Table 1: Storage ───────────────────────────────────────────────────── - W = 60 + W = 78 print("Storage (checkpoint blob bytes)") print("=" * W) print( - f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'savings':>8}" + f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} " + f"{'delta+snap':>12} {'savings':>8}" ) print("-" * W) - 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}" - ) + + def _bytes_or_na(v: Any) -> str: + if v is None: + return "n/a" + if v < 0: + return "n/a" + return _fmt_bytes(v) + + def _ms_or_na(v: Any) -> str: + return "n/a" if v is None else f"{v * 1000:.1f}ms" + + for turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt in rows: + if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0: + ratio_str = "n/a" else: ratio = b_bytes / d_bytes if d_bytes else float("inf") - print( - f"{turns:>6} {_approx_tokens(turns):>10} " - f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} {ratio:>7.0f}x" - ) + ratio_str = f"{ratio:.0f}x" + print( + f"{turns:>6} {_approx_tokens(turns):>10} " + f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} " + f"{_bytes_or_na(s_bytes):>12} {ratio_str:>8}" + ) print("=" * W) print() # ── Table 2: Read latency ───────────────────────────────────────────────── print("Read latency (avg of 5 get_state calls = cost per invoke)") print("=" * W) - print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}") + print( + f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} " + f"{'delta+snap':>12}" + ) print("-" * W) - for turns, b_bytes, d_bytes, b_rt, d_rt in rows: + for turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt in rows: print( f"{turns:>6} {_approx_tokens(turns):>10} " - f"{b_rt * 1000:>10.1f}ms {d_rt * 1000:>10.1f}ms" + f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12} " + f"{_ms_or_na(s_rt):>12}" ) print("=" * W) print() print("Legend:") - print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") - print(" delta = DeltaChannel(add_messages) — O(N) storage") + print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") + print(" delta = DeltaChannel(add_messages) — O(N) storage") + print( + f" delta+snap = DeltaChannel(add_messages, snapshot_every={_SNAPSHOT_EVERY})" + " — O(N) storage, bounded read" + ) print()