Remove Checkpoint.writes

- This has been superseded by saving the individual writes of each task through put_writes()
- Removing this speeds up checkpoint operations as it was duplicating data saved elsewhere already
This commit is contained in:
Nuno Campos
2025-05-25 19:06:46 -07:00
parent 4e8fbe4525
commit 936176eb21
15 changed files with 36 additions and 1399 deletions
@@ -181,11 +181,11 @@ class PostgresSaver(BasePostgresSaver):
"checkpoint_id": value["checkpoint_id"],
}
},
self._load_checkpoint(
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
{
**value["checkpoint"],
"channel_values": self._load_blobs(value["channel_values"]),
},
value["metadata"],
(
{
"configurable": {
@@ -277,11 +277,11 @@ class PostgresSaver(BasePostgresSaver):
"checkpoint_id": value["checkpoint_id"],
}
},
self._load_checkpoint(
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
{
**value["checkpoint"],
"channel_values": self._load_blobs(value["channel_values"]),
},
value["metadata"],
(
{
"configurable": {
@@ -361,8 +361,8 @@ class PostgresSaver(BasePostgresSaver):
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
Jsonb(self._dump_checkpoint(copy)),
self._dump_metadata(get_checkpoint_metadata(config, metadata)),
Jsonb(copy),
Jsonb(get_checkpoint_metadata(config, metadata)),
),
)
return next_config
@@ -168,12 +168,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
"checkpoint_id": value["checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
{
**value["checkpoint"],
"channel_values": self._load_blobs(value["channel_values"]),
},
value["metadata"],
(
{
"configurable": {
@@ -245,12 +244,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
"checkpoint_id": value["checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
{
**value["checkpoint"],
"channel_values": self._load_blobs(value["channel_values"]),
},
value["metadata"],
(
{
"configurable": {
@@ -320,8 +318,8 @@ class AsyncPostgresSaver(BasePostgresSaver):
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
Jsonb(self._dump_checkpoint(copy)),
self._dump_metadata(get_checkpoint_metadata(config, metadata)),
Jsonb(copy),
Jsonb(get_checkpoint_metadata(config, metadata)),
),
)
return next_config
@@ -9,11 +9,8 @@ from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import TASKS
MetadataInput = Optional[dict[str, Any]]
@@ -150,7 +147,6 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
jsonplus_serde = JsonPlusSerializer()
supports_pipeline: bool
def _migrate_pending_sends(
@@ -173,19 +169,6 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
else self.get_next_version(None)
)
def _load_checkpoint(
self,
checkpoint: dict[str, Any],
channel_values: list[tuple[bytes, bytes, bytes]],
) -> Checkpoint:
return {
**checkpoint,
"channel_values": self._load_blobs(channel_values),
}
def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
return checkpoint
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
) -> dict[str, Any]:
@@ -261,14 +244,6 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
for idx, (channel, value) in enumerate(writes)
]
def _load_metadata(self, metadata: dict[str, Any]) -> CheckpointMetadata:
return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata))
def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
serialized_metadata = self.jsonplus_serde.dumps(metadata)
# NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
return serialized_metadata.decode().replace("\\u0000", "")
def get_next_version(self, current: Optional[str]) -> str:
if current is None:
current_v = 0