Remove Checkpoint.pending_sends

- Instead store sends in a Topic channel, removing the need to fetch sends as writes against the parent checkpoint
- Remove deprecated/unused functions in langgraph-checkpoint (will require bumping min range for langgraph-checkpoint in langgraph lib)
- Implement migration of old pending sends in langgraph-checkpoint-postgres
- Ensure parent config of `checkpoint_during=False` checkpoints always points to checkpoints that were also saved
This commit is contained in:
Nuno Campos
2025-05-25 19:06:02 -07:00
parent 85522ec6a6
commit 4e8fbe4525
31 changed files with 559 additions and 508 deletions
+1 -2
View File
@@ -51,7 +51,7 @@ read_config = {"configurable": {"thread_id": "1"}}
checkpointer = MemorySaver()
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -73,7 +73,6 @@ checkpoint = {
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
@@ -1,22 +1,18 @@
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import ( # noqa: UP035
Any,
Dict,
Generic,
List,
Literal,
NamedTuple,
Optional,
Tuple,
TypedDict,
TypeVar,
Union,
)
from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import (
@@ -24,14 +20,10 @@ from langgraph.checkpoint.serde.types import (
INTERRUPT,
RESUME,
SCHEDULED,
ChannelProtocol,
SendProtocol,
)
V = TypeVar("V", int, float, str)
PendingWrite = Tuple[str, str, Any]
# Kept for backwards compat, newer versions of LangGraph no longer use this.
LATEST_VERSION = 2
PendingWrite = tuple[str, str, Any]
# Marked as total=False to allow for future expansion.
@@ -65,10 +57,6 @@ class CheckpointMetadata(TypedDict, total=False):
"""
class TaskInfo(TypedDict):
status: Literal["scheduled", "success", "error"]
ChannelVersions = dict[str, Union[str, int, float]]
@@ -96,22 +84,6 @@ class Checkpoint(TypedDict):
This keeps track of the versions of the channels that each node has seen.
Used to determine which nodes to execute next.
"""
pending_sends: List[SendProtocol]
"""List of inputs pushed to nodes but not yet processed.
Cleared by the next checkpoint."""
# Kept for backwards compat, newer versions of LangGraph no longer use this.
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
v=LATEST_VERSION,
id=str(uuid6(clock_seq=-2)),
ts=datetime.now(timezone.utc).isoformat(),
channel_values={},
channel_versions={},
versions_seen={},
pending_sends=[],
)
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
@@ -122,39 +94,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
channel_values=checkpoint["channel_values"].copy(),
channel_versions=checkpoint["channel_versions"].copy(),
versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
pending_sends=checkpoint.get("pending_sends", []).copy(),
)
# Kept for backwards compat, newer versions of LangGraph no longer use this.
def create_checkpoint(
checkpoint: Checkpoint,
channels: Optional[Mapping[str, ChannelProtocol]],
step: int,
*,
id: Optional[str] = None,
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
if channels is None:
values = checkpoint["channel_values"]
else:
values = {}
for k, v in channels.items():
if k not in checkpoint["channel_versions"]:
continue
try:
values[k] = v.checkpoint()
except EmptyChannelError:
pass
return Checkpoint(
v=LATEST_VERSION,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
pending_sends=checkpoint.get("pending_sends", []),
)
@@ -168,34 +107,6 @@ class CheckpointTuple(NamedTuple):
pending_writes: Optional[List[PendingWrite]] = None
CheckpointThreadId = ConfigurableFieldSpec(
id="thread_id",
annotation=str,
name="Thread ID",
description=None,
default="",
is_shared=True,
)
CheckpointNS = ConfigurableFieldSpec(
id="checkpoint_ns",
annotation=str,
name="Checkpoint NS",
description='Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).',
default="",
is_shared=True,
)
CheckpointId = ConfigurableFieldSpec(
id="checkpoint_id",
annotation=Optional[str],
name="Checkpoint ID",
description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.",
default=None,
is_shared=True,
)
class BaseCheckpointSaver(Generic[V]):
"""Base class for creating a graph checkpointer.
@@ -219,15 +130,6 @@ class BaseCheckpointSaver(Generic[V]):
) -> None:
self.serde = maybe_add_typed_methods(serde or self.serde)
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
"""Define the configuration options for the checkpoint saver.
Returns:
list[ConfigurableFieldSpec]: List of configuration field specs.
"""
return [CheckpointThreadId, CheckpointNS, CheckpointId]
def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
"""Fetch a checkpoint using the given configuration.
@@ -258,7 +160,7 @@ class BaseCheckpointSaver(Generic[V]):
self,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -304,7 +206,7 @@ class BaseCheckpointSaver(Generic[V]):
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
@@ -362,7 +264,7 @@ class BaseCheckpointSaver(Generic[V]):
self,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
@@ -409,7 +311,7 @@ class BaseCheckpointSaver(Generic[V]):
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
@@ -437,7 +339,7 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
def get_next_version(self, current: Optional[V], channel: ChannelProtocol) -> V:
def get_next_version(self, current: Optional[V]) -> V:
"""Generate the next version ID for a channel.
Default is to use integer versions, incrementing by 1. If you override, you can use str/int/float versions,
@@ -445,7 +347,6 @@ class BaseCheckpointSaver(Generic[V]):
Args:
current: The current version identifier (int, float, or str).
channel: The channel being versioned.
Returns:
V: The next version identifier, which must be increasing.
@@ -22,7 +22,6 @@ from langgraph.checkpoint.base import (
get_checkpoint_id,
get_checkpoint_metadata,
)
from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
logger = logging.getLogger(__name__)
@@ -150,19 +149,6 @@ class InMemorySaver(
if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
checkpoint, metadata, parent_checkpoint_id = saved
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
if parent_checkpoint_id:
sends = sorted(
(
(*w, k[1])
for k, w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].items()
if w[1] == TASKS
),
key=lambda w: (w[3], w[0], w[4]),
)
else:
sends = []
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
return CheckpointTuple(
config=config,
@@ -171,7 +157,6 @@ class InMemorySaver(
"channel_values": self._load_blobs(
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
),
"pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
},
metadata=self.serde.loads_typed(metadata),
pending_writes=[
@@ -194,22 +179,7 @@ class InMemorySaver(
checkpoint_id = max(checkpoints.keys())
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
if parent_checkpoint_id:
sends = sorted(
(
(*w, k[1])
for k, w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].items()
if w[1] == TASKS
),
key=lambda w: (w[3], w[0], w[4]),
)
else:
sends = []
checkpoint_ = self.serde.loads_typed(checkpoint)
return CheckpointTuple(
config={
"configurable": {
@@ -223,7 +193,6 @@ class InMemorySaver(
"channel_values": self._load_blobs(
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
),
"pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
},
metadata=self.serde.loads_typed(metadata),
pending_writes=[
@@ -316,20 +285,6 @@ class InMemorySaver(
(thread_id, checkpoint_ns, checkpoint_id)
].values()
if parent_checkpoint_id:
sends = sorted(
(
(*w, k[1])
for k, w in self.writes[
(thread_id, checkpoint_ns, parent_checkpoint_id)
].items()
if w[1] == TASKS
),
key=lambda w: (w[3], w[0], w[4]),
)
else:
sends = []
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
yield CheckpointTuple(
@@ -347,9 +302,6 @@ class InMemorySaver(
checkpoint_ns,
checkpoint_["channel_versions"],
),
"pending_sends": [
self.serde.loads_typed(s[2]) for s in sends
],
},
metadata=metadata,
parent_config=(
@@ -390,7 +342,6 @@ class InMemorySaver(
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
"""
c = checkpoint.copy()
c.pop("pending_sends") # type: ignore[misc]
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc]
@@ -561,7 +512,7 @@ class InMemorySaver(
"""
return self.delete_thread(thread_id)
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
def get_next_version(self, current: Optional[str]) -> str:
if current is None:
current_v = 0
elif isinstance(current, int):
@@ -1,13 +1,4 @@
from collections.abc import Sequence
from typing import (
Any,
Optional,
Protocol,
TypeVar,
runtime_checkable,
)
from typing_extensions import Self
from typing import Any, Protocol, TypeVar, runtime_checkable
ERROR = "__error__"
SCHEDULED = "__scheduled__"
@@ -20,25 +11,6 @@ Update = TypeVar("Update", contravariant=True)
C = TypeVar("C")
class ChannelProtocol(Protocol[Value, Update, C]):
# Mirrors langgraph.channels.base.BaseChannel
@property
def ValueType(self) -> Any: ...
@property
def UpdateType(self) -> Any: ...
def checkpoint(self) -> Optional[C]: ...
def from_checkpoint(self, checkpoint: Optional[C]) -> Self: ...
def update(self, values: Sequence[Update]) -> bool: ...
def get(self) -> Value: ...
def consume(self) -> bool: ...
@runtime_checkable
class SendProtocol(Protocol):
# Mirrors langgraph.constants.Send
+51
View File
@@ -0,0 +1,51 @@
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Any, Optional, Protocol
from langgraph.checkpoint.base import Checkpoint, EmptyChannelError
from langgraph.checkpoint.base.id import uuid6
class ChannelProtocol(Protocol):
def checkpoint(self) -> Optional[Any]: ...
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
v=1,
id=str(uuid6(clock_seq=-2)),
ts=datetime.now(timezone.utc).isoformat(),
channel_values={},
channel_versions={},
versions_seen={},
)
def create_checkpoint(
checkpoint: Checkpoint,
channels: Optional[Mapping[str, ChannelProtocol]],
step: int,
*,
id: Optional[str] = None,
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
if channels is None:
values = checkpoint["channel_values"]
else:
values = {}
for k, v in channels.items():
if k not in checkpoint["channel_versions"]:
continue
try:
values[k] = v.checkpoint()
except EmptyChannelError:
pass
return Checkpoint(
v=1,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
)
+3 -1
View File
@@ -6,10 +6,12 @@ from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
)
from langgraph.checkpoint.memory import InMemorySaver
from tests.checkpoint_utils import ( # type: ignore[import-untyped]
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.memory import InMemorySaver
class TestMemorySaver: