WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-25 09:02:25 +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]:
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+54 -4
View File
@@ -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
)
libs/checkpoint/langgraph/checkpoint/serde/types.py
+1
View File
@@ -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")