From 9b6100bc4fbb6b03e6862e8a9926b0a3bcd75681 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 27 Aug 2024 11:10:18 -0700 Subject: [PATCH] Avoid storing Sends twice in memory and postgres checkpointers - Sends are stored through put_writes, so we don't need to also store them inside checkpoint object - On reading checkpoint, reconstruct pending_sends from the stored writes --- .../langgraph/checkpoint/postgres/__init__.py | 18 +++--- .../langgraph/checkpoint/postgres/aio.py | 26 ++++----- .../langgraph/checkpoint/postgres/base.py | 43 ++++++++------ .../langgraph/checkpoint/memory/__init__.py | 58 +++++++++++++++++-- .../langgraph/checkpoint/serde/types.py | 1 + 5 files changed, 102 insertions(+), 44 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 3ce8e880c..1c064415a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -154,10 +154,11 @@ class PostgresSaver(BasePostgresSaver): "checkpoint_id": value["checkpoint_id"], } }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": self._load_blobs(value["channel_values"]), - }, + self._load_checkpoint( + value["checkpoint"], + value["channel_values"], + value["pending_sends"], + ), self._load_metadata(value["metadata"]), { "configurable": { @@ -232,10 +233,11 @@ class PostgresSaver(BasePostgresSaver): "checkpoint_id": value["checkpoint_id"], } }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": self._load_blobs(value["channel_values"]), - }, + self._load_checkpoint( + value["checkpoint"], + value["channel_values"], + value["pending_sends"], + ), self._load_metadata(value["metadata"]), { "configurable": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 0db61a443..63a2038a3 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver): pipeline (bool): whether to use AsyncPipeline Returns: - PostgresSaver: A new PostgresSaver instance. + AsyncPostgresSaver: A new AsyncPostgresSaver instance. """ async with await AsyncConnection.connect( conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row @@ -137,12 +137,12 @@ class AsyncPostgresSaver(BasePostgresSaver): "checkpoint_id": value["checkpoint_id"], } }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": await asyncio.to_thread( - self._load_blobs, value["channel_values"] - ), - }, + await asyncio.to_thread( + self._load_checkpoint, + value["checkpoint"], + value["channel_values"], + value["pending_sends"], + ), self._load_metadata(value["metadata"]), { "configurable": { @@ -196,12 +196,12 @@ class AsyncPostgresSaver(BasePostgresSaver): "checkpoint_id": value["checkpoint_id"], } }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": await asyncio.to_thread( - self._load_blobs, value["channel_values"] - ), - }, + await asyncio.to_thread( + self._load_checkpoint, + value["checkpoint"], + value["channel_values"], + value["pending_sends"], + ), self._load_metadata(value["metadata"]), { "configurable": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 91b49a162..c0f723f9e 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -1,4 +1,3 @@ -from base64 import b64decode, b64encode from hashlib import md5 from typing import Any, List, Optional, Tuple @@ -13,7 +12,7 @@ from langgraph.checkpoint.base import ( get_checkpoint_id, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.serde.types import ChannelProtocol +from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol MetadataInput = Optional[dict[str, Any]] @@ -58,7 +57,7 @@ MIGRATIONS = [ "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;", ] -SELECT_SQL = """ +SELECT_SQL = f""" select thread_id, checkpoint, @@ -82,7 +81,15 @@ select where cw.thread_id = checkpoints.thread_id and cw.checkpoint_ns = checkpoints.checkpoint_ns and cw.checkpoint_id = checkpoints.checkpoint_id - ) as pending_writes + ) as pending_writes, + ( + select array_agg(array[cw.type::bytea, cw.blob]) + from checkpoint_writes cw + where cw.thread_id = checkpoints.thread_id + and cw.checkpoint_ns = checkpoints.checkpoint_ns + and cw.checkpoint_id = checkpoints.parent_checkpoint_id + and cw.channel = '{TASKS}' + ) as pending_sends from checkpoints """ UPSERT_CHECKPOINT_BLOBS_SQL = """ @@ -116,25 +123,23 @@ class BasePostgresSaver(BaseCheckpointSaver): jsonplus_serde = JsonPlusSerializer() - def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint: - if len(checkpoint["pending_sends"]) == 2 and all( - isinstance(a, str) for a in checkpoint["pending_sends"] - ): - type, bs = checkpoint["pending_sends"] - return { - **checkpoint, - "pending_sends": self.serde.loads_typed((type, b64decode(bs))), - } - - return checkpoint - - def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]: - type, bs = self.serde.dumps_typed(checkpoint["pending_sends"]) + def _load_checkpoint( + self, + checkpoint: dict[str, Any], + channel_values: list[tuple[bytes, bytes, bytes]], + pending_sends: list[tuple[bytes, bytes]], + ) -> Checkpoint: return { **checkpoint, - "pending_sends": (type, b64encode(bs).decode()), + "pending_sends": [ + self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or [] + ], + "channel_values": self._load_blobs(channel_values), } + def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]: + return {**checkpoint, "pending_sends": []} + def _load_blobs( self, blob_values: list[tuple[bytes, bytes, bytes]] ) -> dict[str, Any]: diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 6918de87a..03e1f7efa 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -17,6 +17,7 @@ from langgraph.checkpoint.base import ( SerializerProtocol, get_checkpoint_id, ) +from langgraph.checkpoint.serde.types import TASKS class MemorySaver( @@ -108,9 +109,22 @@ class MemorySaver( 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 = [ + w[2] + for w in self.writes[ + (thread_id, checkpoint_ns, parent_checkpoint_id) + ].values() + if w[1] == TASKS + ] + else: + sends = [] return CheckpointTuple( config=config, - checkpoint=self.serde.loads_typed(checkpoint), + checkpoint={ + **self.serde.loads_typed(checkpoint), + "pending_sends": [self.serde.loads_typed(s) for s in sends], + }, metadata=self.serde.loads_typed(metadata), pending_writes=[ (id, c, self.serde.loads_typed(v)) for id, c, v in writes @@ -130,6 +144,23 @@ class MemorySaver( 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: + print( + [ + *self.writes[ + (thread_id, checkpoint_ns, parent_checkpoint_id) + ].values() + ] + ) + sends = [ + w[2] + for w in self.writes[ + (thread_id, checkpoint_ns, parent_checkpoint_id) + ].values() + if w[1] == TASKS + ] + else: + sends = [] return CheckpointTuple( config={ "configurable": { @@ -138,7 +169,10 @@ class MemorySaver( "checkpoint_id": checkpoint_id, } }, - checkpoint=self.serde.loads_typed(checkpoint), + checkpoint={ + **self.serde.loads_typed(checkpoint), + "pending_sends": [self.serde.loads_typed(s) for s in sends], + }, metadata=self.serde.loads_typed(metadata), pending_writes=[ (id, c, self.serde.loads_typed(v)) for id, c, v in writes @@ -225,6 +259,17 @@ class MemorySaver( (thread_id, checkpoint_ns, checkpoint_id) ].values() + if parent_checkpoint_id: + sends = [ + w[2] + for w in self.writes[ + (thread_id, checkpoint_ns, parent_checkpoint_id) + ].values() + if w[1] == TASKS + ] + else: + sends = [] + yield CheckpointTuple( config={ "configurable": { @@ -233,7 +278,10 @@ class MemorySaver( "checkpoint_id": checkpoint_id, } }, - checkpoint=self.serde.loads_typed(checkpoint), + checkpoint={ + **self.serde.loads_typed(checkpoint), + "pending_sends": [self.serde.loads_typed(s) for s in sends], + }, metadata=metadata, parent_config={ "configurable": { @@ -270,12 +318,14 @@ class MemorySaver( Returns: RunnableConfig: The updated config containing the saved checkpoint's timestamp. """ + c = checkpoint.copy() + c.pop("pending_sends") thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"]["checkpoint_ns"] self.storage[thread_id][checkpoint_ns].update( { checkpoint["id"]: ( - self.serde.dumps_typed(checkpoint), + self.serde.dumps_typed(c), self.serde.dumps_typed(metadata), config["configurable"].get("checkpoint_id"), # parent ) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 26f118bbc..fb0055b51 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -13,6 +13,7 @@ from langchain_core.runnables import RunnableConfig from typing_extensions import Self ERROR = "__error__" +TASKS = "__pregel_tasks" Value = TypeVar("Value") Update = TypeVar("Update")