diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 926faec0f..ffc4d1213 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -32,7 +32,8 @@ Conn = _internal.Conn # For backward compatibility class PostgresSaver(BasePostgresSaver): """Checkpointer that stores checkpoints in a Postgres database.""" - lock: threading.Lock + supports_delta_channels: bool = True + lock: threading.RLock def __init__( self, @@ -48,7 +49,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 @@ -430,43 +431,6 @@ class PostgresSaver(BasePostgresSaver): with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - def get_channel_blob( - self, - thread_id: str, - checkpoint_ns: str, - checkpoint_id: str, - channel: str, - ) -> Any: - """Look up a channel blob by checkpoint ID + channel via checkpoint_blobs.""" - with self._cursor() as cur: - cur.execute( - """ - SELECT cb.type, cb.blob - FROM checkpoint_blobs cb - WHERE cb.thread_id = %s - AND cb.checkpoint_ns = %s - AND cb.channel = %s - AND cb.version = ( - SELECT checkpoint->'channel_versions'->>%s - FROM checkpoints - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s - ) - """, - ( - thread_id, - checkpoint_ns, - channel, - channel, - thread_id, - checkpoint_ns, - checkpoint_id, - ), - ) - row = cur.fetchone() - if row is None: - return NotImplemented - return self.serde.loads_typed((row["type"], row["blob"])) - def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: """ Convert a database row into a CheckpointTuple object. @@ -484,6 +448,7 @@ class PostgresSaver(BasePostgresSaver): value["channel_values"], thread_id=value["thread_id"], checkpoint_ns=value["checkpoint_ns"], + checkpoint_id=value["checkpoint_id"], cur=cur, ) return CheckpointTuple( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index cf8ed4210..61ecb8457 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -13,6 +13,8 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, + DeltaChainValue, + DeltaValue, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -32,6 +34,7 @@ Conn = _ainternal.Conn # For backward compatibility class AsyncPostgresSaver(BasePostgresSaver): """Asynchronous checkpointer that stores checkpoints in a Postgres database.""" + supports_delta_channels: bool = True lock: asyncio.Lock def __init__( @@ -391,42 +394,80 @@ class AsyncPostgresSaver(BasePostgresSaver): async with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - async def aget_channel_blob( + async def _aload_delta_chain( self, thread_id: str, checkpoint_ns: str, checkpoint_id: str, channel: str, - ) -> Any: - """Async look up of a channel blob by checkpoint ID + channel name.""" - async with self._cursor() as cur: - await cur.execute( - """ - SELECT cb.type, cb.blob - FROM checkpoint_blobs cb - WHERE cb.thread_id = %s - AND cb.checkpoint_ns = %s - AND cb.channel = %s - AND cb.version = ( - SELECT checkpoint->'channel_versions'->>%s - FROM checkpoints - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s - ) - """, - ( - thread_id, - checkpoint_ns, - channel, - channel, - thread_id, - checkpoint_ns, - checkpoint_id, - ), + cur: Any, + ) -> DeltaChainValue: + """Fetch the full delta chain for a channel in one recursive CTE query (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 + + 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' ) - row = await cur.fetchone() - if row is None: - return NotImplemented - return self.serde.loads_typed((row["type"], row["blob"])) + SELECT DISTINCT ON (version) type, blob + FROM chain + ORDER BY version ASC + """, + ( + 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) async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: """ @@ -442,11 +483,21 @@ class AsyncPostgresSaver(BasePostgresSaver): """ 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 = [ + k.decode() for k, t, _ in blob_values if t.decode() == "delta" + ] + if delta_channels: + async with self._cursor() as cur: + for channel in delta_channels: + channel_values[channel] = await self._aload_delta_chain( + thread_id, checkpoint_ns, checkpoint_id, channel, cur + ) return CheckpointTuple( { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index c372c3436..cf571c884 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -11,6 +11,8 @@ from langgraph.checkpoint.base import ( WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, + DeltaChainValue, + DeltaValue, get_checkpoint_id, ) from langgraph.checkpoint.serde.types import TASKS @@ -190,18 +192,102 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): *, thread_id: str = "", checkpoint_ns: str = "", + checkpoint_id: str = "", cur: Any = None, ) -> 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 != "empty": + 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 + ) return result + def _load_delta_chain( + 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.""" + 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 + + 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 DISTINCT ON (version) type, blob + FROM chain + ORDER BY version ASC + """, + ( + 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) + def _dump_blobs( self, thread_id: str, diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 48a8b98c1..a9060190a 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -36,9 +36,6 @@ class DeltaValue: """Returned by DeltaChannel.checkpoint(). Represents one step's writes.""" delta: list[Any] - prev_checkpoint_id: ( - str | None - ) # ID of checkpoint containing previous blob; None = chain root @dataclasses.dataclass @@ -478,41 +475,15 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def get_channel_blob( - self, - thread_id: str, - checkpoint_ns: str, - checkpoint_id: str, - channel: str, - ) -> Any: - """Look up a single channel blob by checkpoint ID + channel name. + supports_delta_channels: bool = False + """True if this saver assembles DeltaChannel chains inside get_tuple. - Returns NotImplemented if this saver does not support efficient - per-channel-version blob lookup. The pregel layer will fall back to - get_tuple() traversal in that case. - - Savers with a dedicated blob store (InMemorySaver, PostgresSaver) - should override this for O(1) performance. - """ - return NotImplemented - - async def aget_channel_blob( - self, - thread_id: str, - checkpoint_ns: str, - checkpoint_id: str, - channel: str, - ) -> Any: - """Look up a single channel blob by checkpoint ID + channel name (async). - - Returns NotImplemented if this saver does not support efficient - per-channel-version blob lookup. The pregel layer will fall back to - aget_tuple() traversal in that case. - - Savers with a dedicated blob store (InMemorySaver, PostgresSaver) - should override this for O(1) performance. - """ - return NotImplemented + Savers that set this to True (InMemorySaver, PostgresSaver) assemble the + full DeltaChainValue before returning from get_tuple, so the channel's + from_checkpoint method always receives a DeltaChainValue. Savers that + leave this False will return a raw DeltaValue in channel_values, which + causes DeltaChannel.from_checkpoint to raise a clear error. + """ def get_next_version(self, current: V | None, channel: None) -> V: """Generate the next version ID for a channel. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index eb5296cf0..b6c9895b5 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -20,6 +20,8 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, + DeltaChainValue, + DeltaValue, SerializerProtocol, get_checkpoint_id, get_checkpoint_metadata, @@ -120,51 +122,75 @@ class InMemorySaver( ) -> bool | None: return self.stack.__exit__(__exc_type, __exc_value, __traceback) + supports_delta_channels: bool = True + def _load_blobs( - self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions + self, + 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) if kk not in self.blobs: continue vv = self.blobs[kk] - if vv[0] != "empty": + 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 - def get_channel_blob( + def _assemble_delta_chain( self, thread_id: str, checkpoint_ns: str, checkpoint_id: str, channel: str, - ) -> Any: - """Fast-path blob lookup: checkpoint → channel version → blob.""" + ) -> DeltaChainValue: + """Walk the checkpoint parent tree to collect all delta blobs for a channel.""" ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {}) - entry = ns_storage.get(checkpoint_id) - if entry is None: - return NotImplemented - checkpoint = self.serde.loads_typed(entry[0]) - version = checkpoint["channel_versions"].get(channel) - if version is None: - return NotImplemented - kk = (thread_id, checkpoint_ns, channel, version) - if kk not in self.blobs: - return NotImplemented - vv = self.blobs[kk] - if vv[0] == "empty": - return NotImplemented - return self.serde.loads_typed(vv) - - async def aget_channel_blob( - self, - thread_id: str, - checkpoint_ns: str, - checkpoint_id: str, - channel: str, - ) -> Any: - return self.get_channel_blob(thread_id, checkpoint_ns, checkpoint_id, channel) + 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) + 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) def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the in-memory storage. @@ -192,7 +218,10 @@ class InMemorySaver( checkpoint={ **checkpoint_, "channel_values": self._load_blobs( - thread_id, checkpoint_ns, checkpoint_["channel_versions"] + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + checkpoint_id, ), }, metadata=self.serde.loads_typed(metadata), @@ -228,7 +257,10 @@ class InMemorySaver( checkpoint={ **checkpoint_, "channel_values": self._load_blobs( - thread_id, checkpoint_ns, checkpoint_["channel_versions"] + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + checkpoint_id, ), }, metadata=self.serde.loads_typed(metadata), @@ -338,6 +370,7 @@ class InMemorySaver( thread_id, checkpoint_ns, checkpoint_["channel_versions"], + checkpoint_id, ), }, metadata=metadata, diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index dfa8c2748..889f6758f 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -246,7 +246,7 @@ class JsonPlusSerializer(SerializerProtocol): elif isinstance(obj, bytearray): return "bytearray", obj elif _is_delta_value(obj): - return "delta", _msgpack_enc({"d": obj.delta, "c": obj.prev_checkpoint_id}) + return "delta", _msgpack_enc({"d": obj.delta}) else: try: return "msgpack", _msgpack_enc(obj) @@ -275,7 +275,7 @@ class JsonPlusSerializer(SerializerProtocol): raw = ormsgpack.unpackb( data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS ) - return DeltaValue(delta=raw["d"], prev_checkpoint_id=raw.get("c")) + return DeltaValue(delta=raw["d"]) elif self.pickle_fallback and type_ == "pickle": return pickle.loads(data_) else: diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index 1c8549e10..2c10b61c2 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -990,24 +990,9 @@ def test_delta_value_serde_round_trip() -> None: from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer serde = JsonPlusSerializer() - original = DeltaValue( - delta=[{"type": "human", "content": "hi"}], prev_checkpoint_id="abc-123" - ) + original = DeltaValue(delta=[{"type": "human", "content": "hi"}]) type_tag, blob = serde.dumps_typed(original) assert type_tag == "delta" loaded = serde.loads_typed((type_tag, blob)) assert isinstance(loaded, DeltaValue) assert loaded.delta == original.delta - assert loaded.prev_checkpoint_id == "abc-123" - - -def test_delta_value_serde_chain_root() -> None: - from langgraph.checkpoint.base import DeltaValue - from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer - - serde = JsonPlusSerializer() - original = DeltaValue(delta=[], prev_checkpoint_id=None) - type_tag, blob = serde.dumps_typed(original) - loaded = serde.loads_typed((type_tag, blob)) - assert isinstance(loaded, DeltaValue) - assert loaded.prev_checkpoint_id is None diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 99065325d..b11ef5974 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -311,33 +311,68 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None: class TestInMemorySaverDeltaChannel: - def test_get_channel_blob(self) -> None: - """get_channel_blob returns the deserialized blob for a checkpoint+channel.""" - from langgraph.checkpoint.base import DeltaValue, empty_checkpoint + def test_load_blobs_assembles_delta_chain(self) -> None: + """_load_blobs returns DeltaChainValue for delta channels, not raw DeltaValue.""" + from langgraph.checkpoint.base import ( + DeltaChainValue, + DeltaValue, + empty_checkpoint, + ) saver = InMemorySaver() serde = JsonPlusSerializer() thread_id, ns, channel = "t1", "", "messages" - version = "00000000000000000000000000000001.0000000000000000" - delta = DeltaValue(delta=[{"content": "hi"}], prev_checkpoint_id=None) - saver.blobs[(thread_id, ns, channel, version)] = serde.dumps_typed(delta) + v1 = "00000000000000000000000000000001.0000000000000000" + v2 = "00000000000000000000000000000002.0000000000000000" - cp = empty_checkpoint() - cp["id"] = "cp1" - cp["channel_versions"][channel] = version + 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) + + cp1 = empty_checkpoint() + cp1["id"] = "cp1" + cp1["channel_versions"][channel] = v1 + cp2 = empty_checkpoint() + cp2["id"] = "cp2" + cp2["channel_versions"][channel] = v2 saver.storage[thread_id][ns] = { - "cp1": (serde.dumps_typed(cp), serde.dumps_typed({}), None) + "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), + "cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"), } - result = saver.get_channel_blob(thread_id, ns, "cp1", channel) - assert isinstance(result, DeltaValue) - assert result.delta == [{"content": "hi"}] - assert result.prev_checkpoint_id is None + 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_get_channel_blob_missing(self) -> None: - """get_channel_blob returns NotImplemented when checkpoint or channel not found.""" - saver = InMemorySaver() - assert ( - saver.get_channel_blob("t1", "", "no-such-cp", "messages") is NotImplemented + 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, ) + + 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) + } + + 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"}]] diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 7c17c6ad6..97f5e78b9 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -47,7 +47,6 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): "snapshot_every", "_pending", "_base_version", - "_last_checkpoint_id", "_overwritten", "_steps_since_snapshot", ) @@ -74,7 +73,6 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): self.value = [] self._pending: list[Any] = [] self._base_version: str | None = None - self._last_checkpoint_id: str | None = None self._overwritten: bool = False self._steps_since_snapshot: int = 0 @@ -104,7 +102,6 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): new.value = self.value if self.value is MISSING else self.value.copy() new._pending = self._pending[:] new._base_version = self._base_version - new._last_checkpoint_id = self._last_checkpoint_id new._overwritten = self._overwritten new._steps_since_snapshot = self._steps_since_snapshot return new @@ -126,12 +123,10 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): # the right time regardless of how many prior invocations there were. new._steps_since_snapshot = len(checkpoint.deltas) elif isinstance(checkpoint, DeltaValue): - # Should never reach here — the pregel layer assembles DeltaValues - # into DeltaChainValue before calling from_checkpoint. - raise AssertionError( - "DeltaChannel.from_checkpoint received a raw DeltaValue. " - "This is a bug in the pregel layer — chain assembly should have " - "occurred before from_checkpoint was called." + raise ValueError( + f"Channel '{self.key}' uses DeltaChannel but the checkpointer " + "does not support incremental storage (supports_delta_channels=False). " + "Use InMemorySaver or PostgresSaver, or remove DeltaChannel from your schema." ) else: # Backwards compat: plain list from old BinaryOperatorAggregate checkpoint. @@ -189,10 +184,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): # 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[:], - prev_checkpoint_id=None if self._overwritten else self._last_checkpoint_id, - ) + return DeltaValue(delta=self._pending[:]) def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None: if version != self._base_version: @@ -204,6 +196,5 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): else: self._steps_since_snapshot += 1 self._base_version = version - self._last_checkpoint_id = checkpoint_id self._pending = [] self._overwritten = False diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b1c24de2b..af294a4b7 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1082,6 +1082,27 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): CompiledStateGraph: The compiled `StateGraph`. """ checkpointer = ensure_valid_checkpointer(checkpointer) + + # Warn early if DeltaChannel channels are paired with an incompatible checkpointer. + if checkpointer is not None and checkpointer is not False: + from langgraph.channels.delta import DeltaChannel + + delta_keys = [ + k for k, v in self.channels.items() if isinstance(v, DeltaChannel) + ] + if delta_keys and not getattr( + checkpointer, "supports_delta_channels", False + ): + warnings.warn( + f"Channel(s) {delta_keys} use DeltaChannel but " + f"{type(checkpointer).__name__} does not support incremental " + "channel storage (supports_delta_channels=False). " + "Loading the graph will raise a ValueError. " + "Use InMemorySaver or PostgresSaver.", + UserWarning, + stacklevel=2, + ) + serde_allowlist: set[tuple[str, ...]] | None = None if _serde.STRICT_MSGPACK_ENABLED: schema_types: list[type[Any]] = [ diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 199b013a7..8ebaba89f 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -3,15 +3,8 @@ from __future__ import annotations import logging from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any -from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint.base import ( - BaseCheckpointSaver, - Checkpoint, - DeltaChainValue, - DeltaValue, -) +from langgraph.checkpoint.base import Checkpoint from langgraph.checkpoint.base.id import uuid6 from langgraph._internal._typing import MISSING @@ -22,169 +15,6 @@ LATEST_VERSION = 4 logger = logging.getLogger(__name__) -_MISSING_SENTINEL = object() - - -def _assemble_delta_channels( - checkpoint: Checkpoint, - config: RunnableConfig, - checkpointer: BaseCheckpointSaver, -) -> dict[str, Any]: - """Resolve any DeltaValue entries in checkpoint channel_values to DeltaChainValue. - - Returns a dict of only the channels that needed assembly (others are untouched). - Tries get_channel_blob fast-path first; falls back to get_tuple traversal. - """ - thread_id = str(config["configurable"]["thread_id"]) - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - current_checkpoint_id = checkpoint.get("id") - assembled: dict[str, Any] = {} - - for channel, value in checkpoint["channel_values"].items(): - if not isinstance(value, DeltaValue): - continue - - chain_deltas: list[list[Any]] = [] - base: list[Any] | None = None - cursor: DeltaValue = value - # Pre-seed with current checkpoint ID to guard against self-referential chains. - visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set() - - while True: - chain_deltas.append(cursor.delta) - prev_id = cursor.prev_checkpoint_id - if prev_id is None: - break # chain root - if prev_id in visited: - logger.warning( - "DeltaChannel chain cycle at checkpoint %r for channel %r; breaking", - prev_id, - channel, - ) - break - visited.add(prev_id) - - # Fast path: saver has a dedicated blob store. - blob = checkpointer.get_channel_blob( - thread_id, checkpoint_ns, prev_id, channel - ) - if blob is not NotImplemented: - if isinstance(blob, DeltaValue): - cursor = blob - continue - else: - base = blob # plain list = snapshot root - break - - # Fallback: load the full checkpoint and extract channel value. - parent_config: RunnableConfig = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": prev_id, - } - } - parent_tuple = checkpointer.get_tuple(parent_config) - if parent_tuple is None: - logger.warning( - "DeltaChannel chain broken: checkpoint %r not found for channel %r", - prev_id, - channel, - ) - break - prev_val = parent_tuple.checkpoint["channel_values"].get( - channel, _MISSING_SENTINEL - ) - if prev_val is _MISSING_SENTINEL: - break - elif isinstance(prev_val, DeltaValue): - cursor = prev_val - else: - base = prev_val - break - - chain_deltas.reverse() - assembled[channel] = DeltaChainValue(base=base, deltas=chain_deltas) - - return assembled - - -async def _aassemble_delta_channels( - checkpoint: Checkpoint, - config: RunnableConfig, - checkpointer: BaseCheckpointSaver, -) -> dict[str, Any]: - """Async version of _assemble_delta_channels.""" - thread_id = str(config["configurable"]["thread_id"]) - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - current_checkpoint_id = checkpoint.get("id") - assembled: dict[str, Any] = {} - - for channel, value in checkpoint["channel_values"].items(): - if not isinstance(value, DeltaValue): - continue - - chain_deltas: list[list[Any]] = [] - base: list[Any] | None = None - cursor: DeltaValue = value - visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set() - - while True: - chain_deltas.append(cursor.delta) - prev_id = cursor.prev_checkpoint_id - if prev_id is None: - break - if prev_id in visited: - logger.warning( - "DeltaChannel chain cycle at checkpoint %r for channel %r; breaking", - prev_id, - channel, - ) - break - visited.add(prev_id) - - blob = await checkpointer.aget_channel_blob( - thread_id, checkpoint_ns, prev_id, channel - ) - if blob is not NotImplemented: - if isinstance(blob, DeltaValue): - cursor = blob - continue - else: - base = blob - break - - parent_config: RunnableConfig = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": prev_id, - } - } - parent_tuple = await checkpointer.aget_tuple(parent_config) - if parent_tuple is None: - logger.warning( - "DeltaChannel chain broken: checkpoint %r not found for channel %r", - prev_id, - channel, - ) - break - prev_val = parent_tuple.checkpoint["channel_values"].get( - channel, _MISSING_SENTINEL - ) - if prev_val is _MISSING_SENTINEL: - break - elif isinstance(prev_val, DeltaValue): - cursor = prev_val - else: - base = prev_val - break - - chain_deltas.reverse() - assembled[channel] = DeltaChainValue(base=base, deltas=chain_deltas) - - return assembled - def empty_checkpoint() -> Checkpoint: return Checkpoint( diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 0b4d75008..66365d2d9 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -92,8 +92,6 @@ from langgraph.pregel._algo import ( task_path_str, ) from langgraph.pregel._checkpoint import ( - _aassemble_delta_channels, - _assemble_delta_channels, channels_from_checkpoint, copy_checkpoint, create_checkpoint, @@ -1270,19 +1268,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): else [] ) self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) - # Assemble any DeltaChannel chains before constructing channel objects. - if self.checkpointer is not None: - assembled = _assemble_delta_channels( - self.checkpoint, self.checkpoint_config, self.checkpointer - ) - if assembled: - self.checkpoint = { - **self.checkpoint, - "channel_values": { - **self.checkpoint["channel_values"], - **assembled, - }, - } self.channels, self.managed = channels_from_checkpoint( self.specs, self.checkpoint ) @@ -1487,18 +1472,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): self.submit = await self.stack.enter_async_context( AsyncBackgroundExecutor(self.config) ) - if self.checkpointer is not None: - assembled = await _aassemble_delta_channels( - self.checkpoint, self.checkpoint_config, self.checkpointer - ) - if assembled: - self.checkpoint = { - **self.checkpoint, - "channel_values": { - **self.checkpoint["channel_values"], - **assembled, - }, - } self.channels, self.managed = channels_from_checkpoint( self.specs, self.checkpoint ) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 320c7c51a..46beb59c3 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -122,8 +122,6 @@ from langgraph.pregel._algo import ( ) from langgraph.pregel._call import identifier from langgraph.pregel._checkpoint import ( - _aassemble_delta_channels, - _assemble_delta_channels, channels_from_checkpoint, copy_checkpoint, create_checkpoint, @@ -1052,15 +1050,6 @@ class Pregel( step = saved.metadata.get("step", -1) + 1 stop = step + 2 checkpoint = saved.checkpoint - if isinstance(self.checkpointer, BaseCheckpointSaver): - assembled = _assemble_delta_channels( - checkpoint, saved.config, self.checkpointer - ) - if assembled: - checkpoint = { - **checkpoint, - "channel_values": {**checkpoint["channel_values"], **assembled}, - } channels, managed = channels_from_checkpoint( self.channels, checkpoint, @@ -1181,15 +1170,6 @@ class Pregel( step = saved.metadata.get("step", -1) + 1 stop = step + 2 checkpoint = saved.checkpoint - if isinstance(self.checkpointer, BaseCheckpointSaver): - assembled = await _aassemble_delta_channels( - checkpoint, saved.config, self.checkpointer - ) - if assembled: - checkpoint = { - **checkpoint, - "channel_values": {**checkpoint["channel_values"], **assembled}, - } channels, managed = channels_from_checkpoint( self.channels, checkpoint, @@ -1543,18 +1523,6 @@ class Pregel( if saved is not None: self._migrate_checkpoint(saved.checkpoint) base_checkpoint = saved.checkpoint if saved else empty_checkpoint() - if saved: - assembled = _assemble_delta_channels( - base_checkpoint, saved.config, checkpointer - ) - if assembled: - base_checkpoint = { - **base_checkpoint, - "channel_values": { - **base_checkpoint["channel_values"], - **assembled, - }, - } checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint checkpoint_previous_versions = ( saved.checkpoint["channel_versions"].copy() if saved else {} @@ -2000,18 +1968,6 @@ class Pregel( if saved is not None: self._migrate_checkpoint(saved.checkpoint) base_checkpoint = saved.checkpoint if saved else empty_checkpoint() - if saved: - assembled = await _aassemble_delta_channels( - base_checkpoint, saved.config, checkpointer - ) - if assembled: - base_checkpoint = { - **base_checkpoint, - "channel_values": { - **base_checkpoint["channel_values"], - **assembled, - }, - } checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint checkpoint_previous_versions = ( saved.checkpoint["channel_versions"].copy() if saved else {} diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 3cc328c11..54024ef46 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -134,13 +134,11 @@ def test_delta_channel_basic_two_steps() -> None: d1 = ch.checkpoint() assert isinstance(d1, DeltaValue) assert len(d1.delta) == 1 - assert d1.prev_checkpoint_id is None # first ever step ch.after_checkpoint("v1", checkpoint_id="cid1") # Step 2: another message ch.update([AIMessage(content="hello", id="a1")]) d2 = ch.checkpoint() - assert d2.prev_checkpoint_id == "cid1" assert len(d2.delta) == 1 ch.after_checkpoint("v2") @@ -217,76 +215,26 @@ def test_delta_channel_overwrite_resets_chain() -> None: ch.update([HumanMessage(content="old", id="h1")]) ch.after_checkpoint("v1") - # Overwrite should create a root blob (prev_checkpoint_id=None) ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) d = ch.checkpoint() assert isinstance(d, DeltaValue) - assert d.prev_checkpoint_id is None # chain root assert len(d.delta) == 1 assert d.delta[0].content == "new" + # _overwritten flag must be set so next checkpoint acts as a chain root + assert ch._overwritten is True -def test_delta_channel_assembly_fallback_via_get_tuple() -> None: - """Assembly falls back to get_tuple for savers without get_channel_blob.""" - from unittest.mock import MagicMock - - from langgraph.checkpoint.base import ( - CheckpointTuple, - DeltaChainValue, - DeltaValue, - empty_checkpoint, - ) +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 - from langgraph.pregel._checkpoint import _assemble_delta_channels - - msg1 = {"type": "human", "content": "hello"} - msg2 = {"type": "ai", "content": "world"} - - cp1 = empty_checkpoint() - cp1["id"] = "cp1" - cp1["channel_values"]["messages"] = [msg1] - - cp2 = empty_checkpoint() - cp2["id"] = "cp2" - cp2["channel_values"]["messages"] = DeltaValue( - delta=[msg2], prev_checkpoint_id="cp1" - ) - - saver = MagicMock() - saver.get_channel_blob.return_value = NotImplemented - saver.get_tuple.return_value = CheckpointTuple( - config={ - "configurable": { - "thread_id": "t1", - "checkpoint_ns": "", - "checkpoint_id": "cp1", - } - }, - checkpoint=cp1, - metadata={}, - parent_config=None, - pending_writes=[], - ) - - config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}} - assembled = _assemble_delta_channels(cp2, config, saver) - - assert "messages" in assembled - chain = assembled["messages"] - assert isinstance(chain, DeltaChainValue) - assert chain.base == [msg1] - assert chain.deltas == [[msg2]] - - from langchain_core.messages import AIMessage, HumanMessage spec = DeltaChannel(add_messages) - ch = spec.from_checkpoint(chain) - result = ch.get() - assert len(result) == 2 - assert isinstance(result[0], HumanMessage) and result[0].content == "hello" - assert isinstance(result[1], AIMessage) and result[1].content == "world" + raw = DeltaValue(delta=[{"type": "human", "content": "hello"}]) + with pytest.raises(ValueError, match="supports_delta_channels"): + spec.from_checkpoint(raw) def test_delta_channel_remove_message_delta_and_replay() -> None: @@ -316,7 +264,6 @@ def test_delta_channel_remove_message_delta_and_replay() -> None: ch.update([RemoveMessage(id="a1")]) d2 = ch.checkpoint() assert isinstance(d2, DeltaValue) - assert d2.prev_checkpoint_id == "cid1" assert any(isinstance(w, RemoveMessage) for w in d2.delta) ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == [HumanMessage(content="hi", id="h1")] @@ -349,7 +296,6 @@ def test_delta_channel_update_by_id_delta_and_replay() -> None: ch.update([HumanMessage(content="updated", id="h1")]) d2 = ch.checkpoint() assert isinstance(d2, DeltaValue) - assert d2.prev_checkpoint_id == "cid1" ch.after_checkpoint("v2", checkpoint_id="cid2") assert ch.get() == [HumanMessage(content="updated", id="h1")] @@ -392,7 +338,6 @@ def test_delta_channel_snapshot_every_emits_plain_list() -> None: ch.update([HumanMessage(content="post", id="hpost")]) post = ch.checkpoint() assert isinstance(post, DeltaValue) - assert post.prev_checkpoint_id == "cidsnap" def test_delta_channel_snapshot_every_end_to_end() -> None: @@ -436,60 +381,48 @@ def test_delta_channel_snapshot_every_end_to_end() -> None: assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}" -def test_delta_channel_assembly_fast_path_returns_delta_value() -> None: - """get_channel_blob returning a DeltaValue continues chain traversal (fast-path).""" - from unittest.mock import MagicMock +def test_delta_channel_inmemory_saver_assembles_chain() -> None: + """InMemorySaver assembles the delta chain inside get_tuple (no pregel involvement).""" + from typing import Annotated - from langgraph.checkpoint.base import ( - DeltaChainValue, - DeltaValue, - empty_checkpoint, - ) + 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 - from langgraph.pregel._checkpoint import _assemble_delta_channels - msg1 = {"type": "human", "content": "one"} - msg2 = {"type": "ai", "content": "two"} - msg3 = {"type": "human", "content": "three"} + class State(TypedDict): + messages: Annotated[list, DeltaChannel(add_messages)] - # cp3 → cp2 (DeltaValue) → cp1 (base list) - dv_cp2 = DeltaValue(delta=[msg2], prev_checkpoint_id="cp1") - cp3 = empty_checkpoint() - cp3["id"] = "cp3" - cp3["channel_values"]["messages"] = DeltaValue( - delta=[msg3], prev_checkpoint_id="cp2" - ) + n = {"v": 0} - saver = MagicMock() + def respond(state: State) -> dict: + n["v"] += 1 + return {"messages": [AIMessage(content=f"ok{n['v']}", id=f"ai{n['v']}")]} - def _get_blob(thread_id, ns, checkpoint_id, channel): - if checkpoint_id == "cp2": - return dv_cp2 # DeltaValue — chain continues - if checkpoint_id == "cp1": - return [msg1] # plain list — chain root - return NotImplemented + 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"}} - saver.get_channel_blob.side_effect = _get_blob + graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config) + graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config) - config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}} - assembled = _assemble_delta_channels(cp3, config, saver) + # get_tuple must return a fully assembled DeltaChainValue, not raw DeltaValue + from langgraph.checkpoint.base import DeltaChainValue, DeltaValue - chain = assembled["messages"] - assert isinstance(chain, DeltaChainValue) - assert chain.base == [msg1] - assert chain.deltas == [[msg2], [msg3]] + 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) - spec = DeltaChannel(add_messages) - ch = spec.from_checkpoint(chain) - # add_messages converts dicts to message objects; check by type and content - - result = ch.get() - assert len(result) == 3 - assert result[0].content == "one" - assert result[1].content == "two" - assert result[2].content == "three" + state = graph.get_state(config) + assert len(state.values["messages"]) == 4 # 2 human + 2 AI def test_delta_channel_dict_reducer_fresh_channel() -> None: @@ -526,7 +459,6 @@ def test_delta_channel_dict_reducer_basic_updates() -> None: ch.update([{"b": 2}]) d2 = ch.checkpoint() assert d2.delta == [{"b": 2}] - assert d2.prev_checkpoint_id == "cid1" ch.after_checkpoint("v2") assert ch.get() == {"a": 1, "b": 2} @@ -593,33 +525,43 @@ def test_delta_channel_dict_reducer_with_deletions() -> None: assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"} -def test_delta_channel_assembly_broken_chain_logs_warning() -> None: - """If a prev_checkpoint_id points to a missing checkpoint, log a warning and use partial chain.""" - from unittest.mock import MagicMock +def test_delta_channel_compile_warns_on_incompatible_saver() -> None: + """compile() warns when DeltaChannel is used with a saver that lacks supports_delta_channels.""" + import warnings + from typing import Annotated - from langgraph.checkpoint.base import DeltaValue, empty_checkpoint + from langgraph.checkpoint.base import BaseCheckpointSaver + from typing_extensions import TypedDict - from langgraph.pregel._checkpoint import _assemble_delta_channels + from langgraph.channels.delta import DeltaChannel + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages - cp = empty_checkpoint() - cp["id"] = "cp2" - cp["channel_values"]["messages"] = DeltaValue( - delta=["msg2"], prev_checkpoint_id="cp-missing" - ) + class FakeSaver(BaseCheckpointSaver): + # Third-party saver that has not opted into delta channel support. + supports_delta_channels = False - saver = MagicMock() - saver.get_channel_blob.return_value = NotImplemented - saver.get_tuple.return_value = None # checkpoint not found + def get_tuple(self, config): + return None - config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}} + def put(self, config, checkpoint, metadata, new_versions): + return config - assembled = _assemble_delta_channels(cp, config, saver) + def put_writes(self, config, writes, task_id, task_path=""): + pass - # Should still assemble — with partial chain (just the current delta, base=None) - assert "messages" in assembled - from langgraph.checkpoint.base import DeltaChainValue + def list(self, config, *, filter=None, before=None, limit=None): + return iter([]) - chain = assembled["messages"] - assert isinstance(chain, DeltaChainValue) - assert chain.base is None - assert chain.deltas == [["msg2"]] + class State(TypedDict): + messages: Annotated[list, DeltaChannel(add_messages)] + + builder = StateGraph(State) + builder.add_node("echo", lambda s: {}) + builder.add_edge(START, "echo") + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + builder.compile(checkpointer=FakeSaver()) # type: ignore[arg-type] + + assert any("supports_delta_channels" in str(warning.message) for warning in w) diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index fc6083edb..ee4f2535f 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -16,6 +16,8 @@ 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 @@ -31,6 +33,16 @@ try: 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 + SNAPSHOT_EVERY = 50 # --------------------------------------------------------------------------- @@ -207,20 +219,14 @@ def _approx_tokens(n_turns: int) -> str: # Benchmark matrix # --------------------------------------------------------------------------- -# Turn counts chosen to span from a short session to a long-running agent conversation. -# Storage and time complexity differences are clearly visible by 500 turns. +# 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 = [50, 100, 200, 500] +TURN_COUNTS = [10, 25, 50, 100] def _checkpointer_factories() -> list[tuple[str, Any]]: """Return (label, context_manager_or_none) pairs for available checkpointers.""" - factories: list[tuple[str, Any]] = [("InMemory", None)] - if _SQLITE_AVAILABLE: - import tempfile - - factories.append(("SQLite", tempfile.NamedTemporaryFile(suffix=".db"))) - return factories + return [("InMemory", None)] def run_benchmark() -> None: @@ -232,9 +238,9 @@ def run_benchmark() -> None: print("(5,000 turns × ~200 tokens/turn ≈ 1M tokens — Claude's full context window)") print() - checkpointers: list[tuple[str, Any]] = [("InMemory (fast-path)", None)] - if _SQLITE_AVAILABLE: - checkpointers.append(("SQLite (get_tuple fallback)", "sqlite")) + checkpointers: list[tuple[str, Any]] = [("InMemory", None)] + if _POSTGRES_AVAILABLE: + checkpointers.append(("Postgres (recursive CTE)", "postgres")) for cp_label, cp_hint in checkpointers: print(f"--- Checkpointer: {cp_label} ---") @@ -249,23 +255,24 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: 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 - W = 120 - print("=" * W) - header = ( - f"{'turns':>6} {'ctx size':>10} " - f"{'add_msgs (bytes)':>18} {'delta (bytes)':>15} {'delta+snap (bytes)':>18} " - f"{'storage saved':>14} " - f"{'read: add_msgs':>14} {'read: delta+snap':>16}" - ) - print(header) - print("-" * W) - - results = [] + rows = [] for turns in TURN_COUNTS: with _make_saver() as saver: b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver) @@ -273,51 +280,65 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: 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)) - # For non-InMemory savers, blob_bytes are unavailable (-1); use read times only - if b_bytes < 0 or s_bytes < 0: - b_bytes_str = "n/a" - d_bytes_str = "n/a" - s_bytes_str = "n/a" - storage_ratio_str = "n/a" + # ── Table 1: Storage ───────────────────────────────────────────────────── + W = 80 + print("Storage (checkpoint blob bytes)") + print("=" * W) + print( + f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'delta+snap':>12} {'savings':>8}" + ) + print("-" * W) + storage_results = [] + for turns, b_bytes, d_bytes, s_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}" + ) else: - storage_ratio = b_bytes / s_bytes if s_bytes else float("inf") - b_bytes_str = _fmt_bytes(b_bytes) - d_bytes_str = _fmt_bytes(d_bytes) - s_bytes_str = _fmt_bytes(s_bytes) - storage_ratio_str = f"{storage_ratio:.1f}x" - results.append((turns, b_bytes, s_bytes, b_rt, s_rt, storage_ratio)) - - print( - f"{turns:>6} {_approx_tokens(turns):>10} " - f"{b_bytes_str:>18} {d_bytes_str:>15} {s_bytes_str:>18} " - f"{storage_ratio_str:>14} " - f"{b_rt * 1000:>12.1f}ms {s_rt * 1000:>14.1f}ms" - ) - + ratio = b_bytes / s_bytes if s_bytes else float("inf") + storage_results.append((turns, b_bytes, s_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"{ratio:>7.0f}x" + ) print("=" * W) print() - if results: - best = results[-1] - turns, b_bytes, s_bytes, b_rt, s_rt, ratio = best - print(f"Key findings at max scale ({turns} turns):") + # ── 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} {'delta+snap':>12}" + ) + print("-" * W) + for turns, b_bytes, d_bytes, s_bytes, b_rt, d_rt, s_rt in rows: print( - f" Storage: {_fmt_bytes(b_bytes)} (add_messages) → {_fmt_bytes(s_bytes)} (DeltaChannel+snapshot) — {ratio:.0f}x reduction" + 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" ) + print("=" * W) + print() + + if storage_results: + best = storage_results[-1] + turns, b_bytes, s_bytes, ratio = best + _, _, _, _, b_rt, _, s_rt = rows[-1] print( - f" Read latency: {b_rt * 1000:.1f}ms (add_messages) vs {s_rt * 1000:.1f}ms (DeltaChannel+snapshot)" + 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" ) print() + print("Legend:") + print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") print( - " add_msgs = Annotated[list, add_messages] — current default, O(N²) storage" + " delta = DeltaChannel(add_messages) — O(N) storage, unbounded chain" ) print( - " delta = DeltaChannel(add_messages) — O(N) storage, unbounded chain at read" - ) - print( - f" delta+snap = DeltaChannel(add_messages, snapshot_every={SNAPSHOT_EVERY}) — O(N) storage, O(1) read depth" + f" delta+snap = DeltaChannel(add_messages, snapshot_every={SNAPSHOT_EVERY}) — O(N) storage, bounded read depth" ) print() @@ -327,13 +348,14 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: # --------------------------------------------------------------------------- +@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 [100, 200]: + for turns in [25, 50]: _, _, b_bytes = _run_turns(turns, BinaryState) _, _, d_bytes = _run_turns(turns, DeltaState) _, _, s_bytes = _run_turns(turns, DeltaSnapshotState)