WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 02:07:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Avoid storing Sends twice in memory and postgres checkpointers

Browse Source 
- 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
This commit is contained in:
Nuno Campos
2024-08-27 11:10:18 -07:00
parent 238b562e78
commit 9b6100bc4f
5 changed files with 102 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+10 -8
View File
@@ -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": {
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+13 -13
View File
@@ -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": {
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+24 -19
View File
@@ -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]: