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
+6 -12
View File
@@ -739,7 +739,6 @@
" 'id': '1f029ca3-1f5b-6704-8004-820c16b69a5a',\n",
" 'channel_versions': {'__start__': '00000000000000000000000000000005.0.5290678567601859', 'messages': '00000000000000000000000000000006.0.3205149138784782', 'branch:to:call_model': '00000000000000000000000000000006.0.14611156755133758'}, 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000004.0.5736472536395331'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000005.0.1410174088651449'}},\n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today?), HumanMessage(content=\"what's my name?\"), AIMessage(content='Your name is Bob.')]},\n",
" 'pending_sends': []\n",
" },\n",
" metadata={\n",
" 'source': 'loop',\n",
@@ -856,7 +855,7 @@
" 'id': '1f029ca3-1f5b-6704-8004-820c16b69a5a', \n",
" 'channel_versions': {'__start__': '00000000000000000000000000000005.0.5290678567601859', 'messages': '00000000000000000000000000000006.0.3205149138784782', 'branch:to:call_model': '00000000000000000000000000000006.0.14611156755133758'}, \n",
" 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000004.0.5736472536395331'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000005.0.1410174088651449'}},\n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\"), AIMessage(content='Your name is Bob.')]}, 'pending_sends': []\n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\"), AIMessage(content='Your name is Bob.')]},\n",
" },\n",
" metadata={'source': 'loop', 'writes': {'call_model': {'messages': AIMessage(content='Your name is Bob.')}}, 'step': 4, 'parents': {}, 'thread_id': '1'}, \n",
" parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1f029ca3-1790-6b0a-8003-baf965b6a38f'}}, \n",
@@ -870,8 +869,7 @@
" 'id': '1f029ca3-1790-6b0a-8003-baf965b6a38f', \n",
" 'channel_versions': {'__start__': '00000000000000000000000000000005.0.5290678567601859', 'messages': '00000000000000000000000000000005.0.7935064215293443', 'branch:to:call_model': '00000000000000000000000000000005.0.1410174088651449'}, \n",
" 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000004.0.5736472536395331'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}}, \n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\")], 'branch:to:call_model': None}, \n",
" 'pending_sends': []\n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?'), HumanMessage(content=\"what's my name?\")], 'branch:to:call_model': None}\n",
" }, \n",
" metadata={'source': 'loop', 'writes': None, 'step': 3, 'parents': {}, 'thread_id': '1'}, \n",
" parent_config={...}, \n",
@@ -885,8 +883,7 @@
" 'id': '1f029ca3-1790-616e-8002-9e021694a0cd', \n",
" 'channel_versions': {'__start__': '00000000000000000000000000000004.0.5736472536395331', 'messages': '00000000000000000000000000000003.0.7056767754077798', 'branch:to:call_model': '00000000000000000000000000000003.0.22059023329132854'}, \n",
" 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}}, \n",
" 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"what's my name?\"}]}, 'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}, \n",
" 'pending_sends': []\n",
" 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"what's my name?\"}]}, 'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}\n",
" }, \n",
" metadata={'source': 'input', 'writes': {'__start__': {'messages': [{'role': 'user', 'content': \"what's my name?\"}]}}, 'step': 2, 'parents': {}, 'thread_id': '1'}, \n",
" parent_config={...}, \n",
@@ -900,8 +897,7 @@
" 'id': '1f029ca3-178d-6f54-8001-d7b180db0c89', \n",
" 'channel_versions': {'__start__': '00000000000000000000000000000002.0.18673090920108737', 'messages': '00000000000000000000000000000003.0.7056767754077798', 'branch:to:call_model': '00000000000000000000000000000003.0.22059023329132854'}, \n",
" 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}, 'call_model': {'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}}, \n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}, \n",
" 'pending_sends': []\n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\"), AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')]}\n",
" }, \n",
" metadata={'source': 'loop', 'writes': {'call_model': {'messages': AIMessage(content='Hi Bob! How are you doing today? Is there anything I can help you with?')}}, 'step': 1, 'parents': {}, 'thread_id': '1'}, \n",
" parent_config={...}, \n",
@@ -915,8 +911,7 @@
" 'id': '1f029ca3-0874-6612-8000-339f2abc83b1', \n",
" 'channel_versions': {'__start__': '00000000000000000000000000000002.0.18673090920108737', 'messages': '00000000000000000000000000000002.0.30296526818059655', 'branch:to:call_model': '00000000000000000000000000000002.0.9300422176788571'}, \n",
" 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}}, \n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\")], 'branch:to:call_model': None}, \n",
" 'pending_sends': []\n",
" 'channel_values': {'messages': [HumanMessage(content=\"hi! I'm bob\")], 'branch:to:call_model': None}\n",
" }, \n",
" metadata={'source': 'loop', 'writes': None, 'step': 0, 'parents': {}, 'thread_id': '1'}, \n",
" parent_config={...}, \n",
@@ -930,8 +925,7 @@
" 'id': '1f029ca3-0870-6ce2-bfff-1f3f14c3e565', \n",
" 'channel_versions': {'__start__': '00000000000000000000000000000001.0.7040775356287469'}, \n",
" 'versions_seen': {'__input__': {}}, \n",
" 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"hi! I'm bob\"}]}}, \n",
" 'pending_sends': []\n",
" 'channel_values': {'__start__': {'messages': [{'role': 'user', 'content': \"hi! I'm bob\"}]}}\n",
" }, \n",
" metadata={'source': 'input', 'writes': {'__start__': {'messages': [{'role': 'user', 'content': \"hi! I'm bob\"}]}}, 'step': -1, 'parents': {}, 'thread_id': '1'}, \n",
" parent_config=None, \n",
+2 -4
View File
@@ -25,7 +25,7 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
# call .setup() the first time you're using the checkpointer
checkpointer.setup()
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -47,7 +47,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
@@ -67,7 +66,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -89,7 +88,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
@@ -1,4 +1,5 @@
import threading
from collections import defaultdict
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import Any, Optional
@@ -142,8 +143,36 @@ class PostgresSaver(BasePostgresSaver):
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
with self._cursor() as cur:
cur.execute(query, args, binary=True)
for value in cur:
cur.execute(query, args)
values = cur.fetchall()
if not values:
return
# migrate pending sends if necessary
if to_migrate := [
v
for v in values
if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
]:
cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(
values[0]["thread_id"],
[v["parent_checkpoint_id"] for v in to_migrate],
),
)
grouped_by_parent = defaultdict(list)
for value in to_migrate:
grouped_by_parent[value["parent_checkpoint_id"]].append(value)
for sends in cur:
for value in grouped_by_parent[sends["checkpoint_id"]]:
if value["channel_values"] is None:
value["channel_values"] = []
self._migrate_pending_sends(
sends["sends"],
value["checkpoint"],
value["channel_values"],
)
for value in values:
yield CheckpointTuple(
{
"configurable": {
@@ -155,7 +184,6 @@ class PostgresSaver(BasePostgresSaver):
self._load_checkpoint(
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
@@ -221,37 +249,52 @@ class PostgresSaver(BasePostgresSaver):
cur.execute(
self.SELECT_SQL + where,
args,
binary=True,
)
value = cur.fetchone()
if value is None:
return None
for value in cur:
return CheckpointTuple(
# migrate pending sends if necessary
if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(thread_id, [value["parent_checkpoint_id"]]),
)
if sends := cur.fetchone():
if value["channel_values"] is None:
value["channel_values"] = []
self._migrate_pending_sends(
sends["sends"],
value["checkpoint"],
value["channel_values"],
)
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
}
},
self._load_checkpoint(
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
"checkpoint_id": value["parent_checkpoint_id"],
}
},
self._load_checkpoint(
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None
),
self._load_writes(value["pending_writes"]),
)
}
if value["parent_checkpoint_id"]
else None
),
self._load_writes(value["pending_writes"]),
)
def put(
self,
@@ -1,4 +1,5 @@
import asyncio
from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import Any, Optional
@@ -130,7 +131,35 @@ class AsyncPostgresSaver(BasePostgresSaver):
# if we change this to use .stream() we need to make sure to close the cursor
async with self._cursor() as cur:
await cur.execute(query, args, binary=True)
async for value in cur:
values = await cur.fetchall()
if not values:
return
# migrate pending sends if necessary
if to_migrate := [
v
for v in values
if v["checkpoint"]["v"] < 4 and v["parent_checkpoint_id"]
]:
await cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(
values[0]["thread_id"],
[v["parent_checkpoint_id"] for v in to_migrate],
),
)
grouped_by_parent = defaultdict(list)
for value in to_migrate:
grouped_by_parent[value["parent_checkpoint_id"]].append(value)
async for sends in cur:
for value in grouped_by_parent[sends["checkpoint_id"]]:
if value["channel_values"] is None:
value["channel_values"] = []
self._migrate_pending_sends(
sends["sends"],
value["checkpoint"],
value["channel_values"],
)
for value in values:
yield CheckpointTuple(
{
"configurable": {
@@ -143,7 +172,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
@@ -190,36 +218,52 @@ class AsyncPostgresSaver(BasePostgresSaver):
args,
binary=True,
)
value = await cur.fetchone()
if value is None:
return None
async for value in cur:
return CheckpointTuple(
# migrate pending sends if necessary
if value["checkpoint"]["v"] < 4 and value["parent_checkpoint_id"]:
await cur.execute(
self.SELECT_PENDING_SENDS_SQL,
(thread_id, [value["parent_checkpoint_id"]]),
)
if sends := await cur.fetchone():
if value["channel_values"] is None:
value["channel_values"] = []
self._migrate_pending_sends(
sends["sends"],
value["checkpoint"],
value["channel_values"],
)
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
"checkpoint_id": value["parent_checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None
),
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
}
if value["parent_checkpoint_id"]
else None
),
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
async def aput(
self,
@@ -14,7 +14,7 @@ from langgraph.checkpoint.base import (
get_checkpoint_id,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
from langgraph.checkpoint.serde.types import TASKS
MetadataInput = Optional[dict[str, Any]]
@@ -72,7 +72,7 @@ MIGRATIONS = [
"""ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT '';""",
]
SELECT_SQL = f"""
SELECT_SQL = """
select
thread_id,
checkpoint,
@@ -96,17 +96,20 @@ 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,
(
select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx)
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
) as pending_writes
from checkpoints """
SELECT_PENDING_SENDS_SQL = f"""
select
checkpoint_id,
array_agg(array[type::bytea, blob] order by task_path, task_id, idx) as sends
from checkpoint_writes
where thread_id = %s
and checkpoint_id = any(%s)
and channel = '{TASKS}'
group by checkpoint_id
"""
UPSERT_CHECKPOINT_BLOBS_SQL = """
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
VALUES (%s, %s, %s, %s, %s, %s)
@@ -140,6 +143,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
MIGRATIONS = MIGRATIONS
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
@@ -149,22 +153,38 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
jsonplus_serde = JsonPlusSerializer()
supports_pipeline: bool
def _migrate_pending_sends(
self,
pending_sends: list[tuple[bytes, bytes]],
checkpoint: dict[str, Any],
channel_values: list[tuple[bytes, bytes, bytes]],
) -> None:
if not pending_sends:
return
# add to values
enc, blob = self.serde.dumps_typed(
[self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends],
)
channel_values.append((TASKS.encode(), enc.encode(), blob))
# add to versions
checkpoint["channel_versions"][TASKS] = (
max(checkpoint["channel_versions"].values())
if checkpoint["channel_versions"]
else self.get_next_version(None)
)
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": [
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": []}
return checkpoint
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
@@ -249,7 +269,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
# 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], channel: ChannelProtocol) -> str:
def get_next_version(self, current: Optional[str]) -> str:
if current is None:
current_v = 0
elif isinstance(current, int):
@@ -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"],
)
+51 -2
View File
@@ -14,10 +14,10 @@ from langgraph.checkpoint.base import (
EXCLUDED_METADATA_KEYS,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.serde.types import TASKS
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -263,3 +263,52 @@ async def test_null_chars(saver_name: str, test_data) -> None:
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
0
].metadata["my_key"] == "abc"
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_pending_sends_migration(saver_name: str) -> None:
async with _saver(saver_name) as saver:
config = {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
}
}
# create the first checkpoint
# and put some pending sends
checkpoint_0 = empty_checkpoint()
config = await saver.aput(config, checkpoint_0, {}, {})
await saver.aput_writes(
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
)
await saver.aput_writes(config, [(TASKS, "send-3")], task_id="task-2")
# check that fetching checkpoint_0 doesn't attach pending sends
# (they should be attached to the next checkpoint)
tuple_0 = await saver.aget_tuple(config)
assert tuple_0.checkpoint["channel_values"] == {}
assert tuple_0.checkpoint["channel_versions"] == {}
# create the second checkpoint
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
config = await saver.aput(config, checkpoint_1, {}, {})
# check that pending sends are attached to checkpoint_1
tuple_1 = await saver.aget_tuple(config)
assert tuple_1.checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in tuple_1.checkpoint["channel_versions"]
# check that list also applies the migration
search_results = [
c async for c in saver.alist({"configurable": {"thread_id": "thread-1"}})
]
assert len(search_results) == 2
assert search_results[-1].checkpoint["channel_values"] == {}
assert search_results[-1].checkpoint["channel_versions"] == {}
assert search_results[0].checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in search_results[0].checkpoint["channel_versions"]
+51 -2
View File
@@ -15,10 +15,10 @@ from langgraph.checkpoint.base import (
EXCLUDED_METADATA_KEYS,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.serde.types import TASKS
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -258,3 +258,52 @@ def test_nonnull_migrations() -> None:
for migration in PostgresSaver.MIGRATIONS:
statement = _leading_comment_remover.sub("", migration).split()[0]
assert statement.strip()
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
def test_pending_sends_migration(saver_name: str) -> None:
with _saver(saver_name) as saver:
config = {
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
}
}
# create the first checkpoint
# and put some pending sends
checkpoint_0 = empty_checkpoint()
config = saver.put(config, checkpoint_0, {}, {})
saver.put_writes(
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
)
saver.put_writes(config, [(TASKS, "send-3")], task_id="task-2")
# check that fetching checkpoint_0 doesn't attach pending sends
# (they should be attached to the next checkpoint)
tuple_0 = saver.get_tuple(config)
assert tuple_0.checkpoint["channel_values"] == {}
assert tuple_0.checkpoint["channel_versions"] == {}
# create the second checkpoint
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
config = saver.put(config, checkpoint_1, {}, {})
# check that pending sends are attached to checkpoint_1
checkpoint_1 = saver.get_tuple(config)
assert checkpoint_1.checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in checkpoint_1.checkpoint["channel_versions"]
# check that list also applies the migration
search_results = [
c for c in saver.list({"configurable": {"thread_id": "thread-1"}})
]
assert len(search_results) == 2
assert search_results[-1].checkpoint["channel_values"] == {}
assert search_results[-1].checkpoint["channel_versions"] == {}
assert search_results[0].checkpoint["channel_values"] == {
TASKS: ["send-1", "send-2", "send-3"]
}
assert TASKS in search_results[0].checkpoint["channel_versions"]
+2 -4
View File
@@ -12,7 +12,7 @@ read_config = {"configurable": {"thread_id": "1"}}
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -34,7 +34,6 @@ with SqliteSaver.from_conn_string(":memory:") as checkpointer:
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
@@ -54,7 +53,7 @@ from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
checkpoint = {
"v": 2,
"v": 4,
"ts": "2024-07-31T20:14:19.804150+00:00",
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
"channel_values": {
@@ -76,7 +75,6 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
"start:node": 2
}
},
"pending_sends": [],
}
# store checkpoint
@@ -19,7 +19,6 @@ from langgraph.checkpoint.base import (
get_checkpoint_metadata,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import ChannelProtocol
from langgraph.checkpoint.sqlite.utils import search_where
_AIO_ERROR_MSG = (
@@ -535,14 +534,13 @@ class SqliteSaver(BaseCheckpointSaver[str]):
"""
raise NotImplementedError(_AIO_ERROR_MSG)
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
def get_next_version(self, current: Optional[str]) -> str:
"""Generate the next version ID for a channel.
This method creates a new version identifier for a channel based on its current version.
Args:
current (Optional[str]): The current version identifier of the channel.
channel (BaseChannel): The channel being versioned.
Returns:
str: The next version identifier, which is guaranteed to be monotonically increasing.
@@ -19,7 +19,6 @@ from langgraph.checkpoint.base import (
get_checkpoint_metadata,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import ChannelProtocol
from langgraph.checkpoint.sqlite.utils import search_where
T = TypeVar("T", bound=Callable)
@@ -590,14 +589,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
)
await self.conn.commit()
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
def get_next_version(self, current: Optional[str]) -> str:
"""Generate the next version ID for a channel.
This method creates a new version identifier for a channel based on its current version.
Args:
current (Optional[str]): The current version identifier of the channel.
channel (BaseChannel): The channel being versioned.
Returns:
str: The next version identifier, which is guaranteed to be monotonically increasing.
@@ -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"],
)
@@ -6,10 +6,9 @@ from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
class TestAsyncSqliteSaver:
+1 -2
View File
@@ -6,11 +6,10 @@ from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.utils import _metadata_predicate, search_where
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
class TestSqliteSaver:
+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:
+5 -3
View File
@@ -71,12 +71,14 @@ class Topic(
return empty
def update(self, values: Sequence[Union[Value, list[Value]]]) -> bool:
current = list(self.values)
updated = False
if not self.accumulate:
updated = bool(self.values)
self.values = list[Value]()
if flat_values := flatten(values):
if flat_values := tuple(flatten(values)):
updated = True
self.values.extend(flat_values)
return self.values != current
return updated
def get(self) -> Sequence[Value]:
if self.values:
-1
View File
@@ -122,7 +122,6 @@ RESERVED = {
ERROR,
NO_WRITES,
SCHEDULED,
TASKS,
# reserved config.configurable keys
CONFIG_KEY_SEND,
CONFIG_KEY_READ,
+14 -56
View File
@@ -22,17 +22,12 @@ from langchain_core.runnables.config import (
get_callback_manager_for_config,
)
from langchain_core.runnables.graph import Graph
from langchain_core.runnables.utils import (
ConfigurableFieldSpec,
get_unique_config_specs,
)
from pydantic import BaseModel
from typing_extensions import Self
from langgraph.cache.base import BaseCache
from langgraph.channels.base import (
BaseChannel,
)
from langgraph.channels.base import BaseChannel
from langgraph.channels.topic import Topic
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
@@ -49,7 +44,6 @@ from langgraph.constants import (
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_NODE_FINISHED,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
CONFIG_KEY_RUNNER_SUBMIT,
CONFIG_KEY_SEND,
CONFIG_KEY_STORE,
@@ -66,6 +60,7 @@ from langgraph.constants import (
NULL_TASK_ID,
PUSH,
SCHEDULED,
TASKS,
)
from langgraph.errors import (
ErrorCode,
@@ -104,6 +99,7 @@ from langgraph.types import (
CachePolicy,
Checkpointer,
Interrupt,
Send,
StateSnapshot,
StateUpdate,
StreamChunk,
@@ -117,8 +113,7 @@ from langgraph.utils.config import (
patch_configurable,
recast_checkpoint_ns,
)
from langgraph.utils.fields import get_enhanced_type_hints
from langgraph.utils.pydantic import create_model, is_supported_by_pydantic
from langgraph.utils.pydantic import create_model
from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined]
from langgraph.utils.runnable import (
Runnable,
@@ -706,6 +701,12 @@ class Pregel(PregelProtocol):
k: v.build() if isinstance(v, NodeBuilder) else v for k, v in nodes.items()
}
self.channels = channels or {}
if TASKS in self.channels and not isinstance(self.channels[TASKS], Topic):
raise ValueError(
f"Channel '{TASKS}' is reserved and cannot be used in the graph."
)
else:
self.channels[TASKS] = Topic(Send, accumulate=False)
self.stream_mode = stream_mode
self.stream_eager = stream_eager
self.output_channels = output_channels
@@ -830,54 +831,10 @@ class Pregel(PregelProtocol):
self.trigger_to_nodes = _trigger_to_nodes(self.nodes)
return self
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
return [
spec
for spec in get_unique_config_specs(
[spec for node in self.nodes.values() for spec in node.config_specs]
+ (
self.checkpointer.config_specs
if isinstance(self.checkpointer, BaseCheckpointSaver)
else []
)
+ (
[
ConfigurableFieldSpec(
id=name,
annotation=typ,
default=default,
description=description,
)
for name, typ, default, description in get_enhanced_type_hints(
self.config_type
)
]
if self.config_type is not None
else []
)
)
# these are provided by the Pregel class
if spec.id
not in [
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_RESUMING,
]
]
def config_schema(self, *, include: Sequence[str] | None = None) -> type[BaseModel]:
# If the config type is not set explicitly, we will try to infer it.
# If the config type is provided, but isn't directly supported by pydantic
# (e.g., vanilla python class), we will also delegate to the parent class,
# which handles cases where Pydantic doesn't support the type.
if self.config_type is None or not is_supported_by_pydantic(self.config_type):
return super().config_schema(include=include)
include = include or []
fields = {
"configurable": (self.config_type, None),
**({"configurable": (self.config_type, None)} if self.config_type else {}),
**{
field_name: (field_type, None)
for field_name, field_type in get_type_hints(RunnableConfig).items()
@@ -2359,7 +2316,8 @@ class Pregel(PregelProtocol):
checkpointer = self.checkpointer
if checkpointer and not config.get(CONF):
raise ValueError(
f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in checkpointer.config_specs]}"
"Checkpointer requires one or more of the following 'configurable' "
"keys: thread_id, checkpoint_ns, checkpoint_id"
)
if CONFIG_KEY_STORE in config.get(CONF, {}):
store: BaseStore | None = config[CONF][CONFIG_KEY_STORE]
+51 -79
View File
@@ -25,6 +25,7 @@ from langchain_core.runnables.config import RunnableConfig
from xxhash import xxh3_128_hexdigest
from langgraph.channels.base import BaseChannel
from langgraph.channels.topic import Topic
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
ChannelVersions,
@@ -81,7 +82,7 @@ from langgraph.types import (
)
from langgraph.utils.config import merge_configs, patch_config
GetNextVersion = Callable[[Optional[V], BaseChannel], V]
GetNextVersion = Callable[[Optional[V]], V]
SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
@@ -211,7 +212,7 @@ def local_read(
return values
def increment(current: Optional[int], channel: BaseChannel) -> int:
def increment(current: Optional[int]) -> int:
"""Default channel versioning function, increments the current int version."""
return current + 1 if current is not None else 1
@@ -232,6 +233,7 @@ def apply_writes(
channels: The channels to update.
tasks: The tasks to apply writes from.
get_next_version: Optional function to determine the next version of a channel.
trigger_to_nodes: Mapping of channel names to the set of nodes that can be triggered by updates to that channel.
Returns:
Set of channels that were updated in this step.
@@ -255,10 +257,14 @@ def apply_writes(
)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
if get_next_version is None:
next_version = None
else:
max_version = None
next_version = get_next_version(
max(checkpoint["channel_versions"].values())
if checkpoint["channel_versions"]
else None
)
# Consume all channels that were read
for chan in {
@@ -267,15 +273,8 @@ def apply_writes(
for chan in task.triggers
if chan not in RESERVED and chan in channels
}:
if channels[chan].consume() and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version,
channels[chan],
)
# clear pending sends
if checkpoint["pending_sends"] and bump_step:
checkpoint["pending_sends"].clear()
if channels[chan].consume() and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# Group writes by channel
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
@@ -283,8 +282,6 @@ def apply_writes(
for chan, val in task.writes:
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
pass
elif chan == TASKS:
checkpoint["pending_sends"].append(val)
elif chan in channels:
pending_writes_by_channel[chan].append(val)
else:
@@ -292,21 +289,12 @@ def apply_writes(
f"Task {task.name} with path {task.path} wrote to unknown channel {chan}, ignoring it."
)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = None
# Apply writes to channels
updated_channels: set[str] = set()
for chan, vals in pending_writes_by_channel.items():
if chan in channels:
if channels[chan].update(vals) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version,
channels[chan],
)
if channels[chan].update(vals) and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
updated_channels.add(chan)
@@ -315,27 +303,17 @@ def apply_writes(
if bump_step:
for chan in channels:
if channels[chan].is_available() and chan not in updated_channels:
if channels[chan].update(EMPTY_SEQ) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version,
channels[chan],
)
if channels[chan].update(EMPTY_SEQ) and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
updated_channels.add(chan)
# If this is (tentatively) the last superstep, notify all channels of finish
if (
bump_step
and not checkpoint["pending_sends"]
and updated_channels.isdisjoint(trigger_to_nodes)
):
if bump_step and updated_channels.isdisjoint(trigger_to_nodes):
for chan in channels:
if channels[chan].finish() and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version,
channels[chan],
)
if channels[chan].finish() and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
updated_channels.add(chan)
@@ -344,17 +322,6 @@ def apply_writes(
return updated_channels
def has_next_tasks(
trigger_to_nodes: Mapping[str, Sequence[str]],
updated_channels: set[str],
checkpoint: Checkpoint,
) -> bool:
"""Check if there are any tasks that should be run in the next step."""
return bool(checkpoint["pending_sends"]) or not updated_channels.isdisjoint(
trigger_to_nodes
)
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
@@ -448,30 +415,32 @@ def prepare_next_tasks(
checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", ""))
null_version = checkpoint_null_version(checkpoint)
tasks: list[Union[PregelTask, PregelExecutableTask]] = []
# Consume pending_sends from previous step
for idx, _ in enumerate(checkpoint["pending_sends"]):
if task := prepare_single_task(
(PUSH, idx),
None,
checkpoint=checkpoint,
checkpoint_id_bytes=checkpoint_id_bytes,
checkpoint_null_version=null_version,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
stop=stop,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
input_cache=input_cache,
cache_policy=cache_policy,
retry_policy=retry_policy,
):
tasks.append(task)
# Consume pending tasks
tasks_channel = cast(Optional[Topic[Send]], channels.get(TASKS))
if tasks_channel and tasks_channel.is_available():
for idx, _ in enumerate(tasks_channel.get()):
if task := prepare_single_task(
(PUSH, idx),
None,
checkpoint=checkpoint,
checkpoint_id_bytes=checkpoint_id_bytes,
checkpoint_null_version=null_version,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
stop=stop,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
input_cache=input_cache,
cache_policy=cache_policy,
retry_policy=retry_policy,
):
tasks.append(task)
# This section is an optimization that allows which nodes will be active
# during the next step.
@@ -656,9 +625,12 @@ def prepare_single_task(
# SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
idx = cast(int, task_path[1])
if idx >= len(checkpoint["pending_sends"]):
if not channels[TASKS].is_available():
return
packet = checkpoint["pending_sends"][idx]
sends: Sequence[Send] = channels[TASKS].get()
if idx < 0 or idx >= len(sends):
return
packet = sends[idx]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending sends"
@@ -8,7 +8,7 @@ from langgraph.checkpoint.base.id import uuid6
from langgraph.constants import MISSING
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
LATEST_VERSION = 3
LATEST_VERSION = 4
def empty_checkpoint() -> Checkpoint:
@@ -19,7 +19,6 @@ def empty_checkpoint() -> Checkpoint:
channel_values={},
channel_versions={},
versions_seen={},
pending_sends=[],
)
@@ -49,7 +48,6 @@ def create_checkpoint(
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
pending_sends=checkpoint.get("pending_sends", []),
)
+1 -7
View File
@@ -65,7 +65,6 @@ from langgraph.constants import (
RESUME,
SCHEDULED,
TAG_HIDDEN,
TASKS,
)
from langgraph.errors import (
CheckpointNotLatest,
@@ -318,9 +317,6 @@ class PregelLoop:
"""Put writes for a task, to be read by the next tick."""
if not writes:
return
# always checkpoint writes containing Send, as they are fetched from the
# parent checkpoint, not the current one
checkpoint_during = self.checkpoint_during or any(w[0] == TASKS for w in writes)
# deduplicate writes to special channels, last write wins
if all(w[0] in WRITES_IDX_MAP for w in writes):
writes = list({w[0]: w for w in writes}.values())
@@ -330,7 +326,7 @@ class PregelLoop:
]
# save writes
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
if checkpoint_during and self.checkpointer_put_writes is not None:
if self.checkpoint_during and self.checkpointer_put_writes is not None:
config = patch_configurable(
self.checkpoint_config,
{
@@ -806,7 +802,6 @@ class PregelLoop:
else self.stream_keys
),
)
self.checkpoint_id_prev = self.checkpoint["id"] if self.step > -1 else None
# do checkpoint?
do_checkpoint = self._checkpointer_put_after_previous is not None and (
exiting or self.checkpoint_during
@@ -836,7 +831,6 @@ class PregelLoop:
CONF: {
**self.checkpoint_config[CONF],
# this is guaranteed to be set by code above
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint_id_prev,
CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
CONFIG_KEY_CHECKPOINT_NS, ""
),
+1 -2
View File
@@ -1,7 +1,6 @@
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.constants import PULL, PUSH
from langgraph.pregel.algo import prepare_next_tasks, task_path_str
from langgraph.pregel.checkpoint import channels_from_checkpoint
from langgraph.pregel.checkpoint import channels_from_checkpoint, empty_checkpoint
def test_prepare_next_tasks() -> None:
@@ -350,7 +350,6 @@ SAVED_CHECKPOINTS = {
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -413,7 +412,6 @@ SAVED_CHECKPOINTS = {
"docs": ["doc1", "doc2", "doc3", "doc4"],
"branch:to:qa": None,
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -491,7 +489,6 @@ SAVED_CHECKPOINTS = {
"branch:to:retriever_one": None,
"docs": ["doc3", "doc4"],
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -548,7 +545,6 @@ SAVED_CHECKPOINTS = {
"branch:to:analyzer_one": None,
"branch:to:retriever_two": None,
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -605,7 +601,6 @@ SAVED_CHECKPOINTS = {
"query": "what is weather in sf",
"branch:to:rewrite_query": None,
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -656,7 +651,6 @@ SAVED_CHECKPOINTS = {
},
"versions_seen": {"__input__": {}},
"channel_values": {"__start__": {"query": "what is weather in sf"}},
"pending_sends": [],
},
metadata={
"source": "input",
@@ -744,7 +738,6 @@ SAVED_CHECKPOINTS = {
"answer": "doc1,doc2,doc3,doc4",
"qa": "qa",
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -808,7 +801,6 @@ SAVED_CHECKPOINTS = {
"docs": ["doc1", "doc2", "doc3", "doc4"],
"retriever_one": "retriever_one",
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -889,7 +881,6 @@ SAVED_CHECKPOINTS = {
"docs": ["doc3", "doc4"],
"retriever_two": "retriever_two",
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -950,7 +941,6 @@ SAVED_CHECKPOINTS = {
"rewrite_query": "rewrite_query",
"branch:to:retriever_two": "rewrite_query",
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1012,7 +1002,6 @@ SAVED_CHECKPOINTS = {
"query": "what is weather in sf",
"start:rewrite_query": "__start__",
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1063,7 +1052,6 @@ SAVED_CHECKPOINTS = {
},
"versions_seen": {"__input__": {}},
"channel_values": {"__start__": {"query": "what is weather in sf"}},
"pending_sends": [],
},
metadata={
"source": "input",
@@ -1151,7 +1139,6 @@ SAVED_CHECKPOINTS = {
"retriever_one": "00000000000000000000000000000005.0.222301724202566"
},
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1215,7 +1202,6 @@ SAVED_CHECKPOINTS = {
"analyzer_one": "00000000000000000000000000000004.0.2684613370070208"
},
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1296,7 +1282,6 @@ SAVED_CHECKPOINTS = {
"branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000003.0.8992241767805405"
},
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1357,7 +1342,6 @@ SAVED_CHECKPOINTS = {
"start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445"
},
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1419,7 +1403,6 @@ SAVED_CHECKPOINTS = {
"__start__": "00000000000000000000000000000001.0.6759219622820284"
},
},
"pending_sends": [],
},
metadata={
"source": "loop",
@@ -1470,7 +1453,6 @@ SAVED_CHECKPOINTS = {
"__start__": "00000000000000000000000000000001.0.6759219622820284"
},
"versions_seen": {"__input__": {}},
"pending_sends": [],
},
metadata={
"source": "input",
+4 -2
View File
@@ -7416,8 +7416,10 @@ def test_send_dedupe_on_resume(
if checkpoint_during:
assert history == expected_history
else:
assert history[0] == expected_history[0]
assert history[1] == expected_history[2]
assert history[0] == expected_history[0]._replace(
parent_config=history[1].config
)
assert history[1] == expected_history[2]._replace(parent_config=None)
def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None:
+10 -30
View File
@@ -30,7 +30,6 @@ from syrupy import SnapshotAssertion
from typing_extensions import TypedDict
from langgraph.cache.base import BaseCache
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
@@ -237,7 +236,7 @@ def test_checkpoint_errors() -> None:
raise ValueError("Faulty put_writes")
class FaultyVersionCheckpointer(InMemorySaver):
def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int:
def get_next_version(self, current: Optional[int]) -> int:
raise ValueError("Faulty get_next_version")
def logic(inp: str) -> str:
@@ -1130,10 +1129,9 @@ def test_pending_writes_resume(
}
},
checkpoint={
"v": 3,
"v": 4,
"id": AnyStr(),
"ts": AnyStr(),
"pending_sends": [],
"versions_seen": {
"one": {
"branch:to:one": AnyVersion(),
@@ -1188,10 +1186,9 @@ def test_pending_writes_resume(
}
},
checkpoint={
"v": 3,
"v": 4,
"id": AnyStr(),
"ts": AnyStr(),
"pending_sends": [],
"versions_seen": {
"__input__": {},
"__start__": {
@@ -1223,11 +1220,11 @@ def test_pending_writes_resume(
"checkpoint_ns": "",
"checkpoint_id": (
checkpoints[2].config["configurable"]["checkpoint_id"]
if checkpoint_during
else AnyStr()
),
}
},
}
if checkpoint_during
else None,
pending_writes=(
UnsortedSequence(
(AnyStr(), "value", 2),
@@ -1254,10 +1251,9 @@ def test_pending_writes_resume(
}
},
checkpoint={
"v": 3,
"v": 4,
"id": AnyStr(),
"ts": AnyStr(),
"pending_sends": [],
"versions_seen": {"__input__": {}},
"channel_versions": {
"__start__": AnyVersion(),
@@ -1407,7 +1403,7 @@ def test_imp_task(
) -> None:
mapper_calls = 0
class Configurable:
class Configurable(TypedDict):
model: str
@task()
@@ -1438,25 +1434,9 @@ def test_imp_task(
"$defs": {
"Configurable": {
"properties": {
"model": {"default": None, "title": "Model", "type": "string"},
"checkpoint_id": {
"anyOf": [{"type": "string"}, {"type": "null"}],
"default": None,
"description": "Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.",
"title": "Checkpoint ID",
},
"checkpoint_ns": {
"default": "",
"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).',
"title": "Checkpoint NS",
"type": "string",
},
"thread_id": {
"default": "",
"title": "Thread ID",
"type": "string",
},
"model": {"title": "Model", "type": "string"},
},
"required": ["model"],
"title": "Configurable",
"type": "object",
}
+12 -14
View File
@@ -29,7 +29,6 @@ from syrupy import SnapshotAssertion
from typing_extensions import TypedDict
from langgraph.cache.base import BaseCache
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
@@ -104,7 +103,7 @@ async def test_checkpoint_errors() -> None:
raise ValueError("Faulty put_writes")
class FaultyVersionCheckpointer(InMemorySaver):
def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int:
def get_next_version(self, current: Optional[int]) -> int:
raise ValueError("Faulty get_next_version")
def logic(inp: str) -> str:
@@ -1940,10 +1939,9 @@ async def test_pending_writes_resume(
}
},
checkpoint={
"v": 3,
"v": 4,
"id": AnyStr(),
"ts": AnyStr(),
"pending_sends": [],
"versions_seen": {
"one": {
"branch:to:one": AnyVersion(),
@@ -1998,10 +1996,9 @@ async def test_pending_writes_resume(
}
},
checkpoint={
"v": 3,
"v": 4,
"id": AnyStr(),
"ts": AnyStr(),
"pending_sends": [],
"versions_seen": {
"__input__": {},
"__start__": {
@@ -2031,11 +2028,11 @@ async def test_pending_writes_resume(
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"]
if checkpoint_during
else AnyStr(),
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
}
},
}
if checkpoint_during
else None,
pending_writes=UnsortedSequence(
(AnyStr(), "value", 2),
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
@@ -2060,10 +2057,9 @@ async def test_pending_writes_resume(
}
},
checkpoint={
"v": 3,
"v": 4,
"id": AnyStr(),
"ts": AnyStr(),
"pending_sends": [],
"versions_seen": {"__input__": {}},
"channel_versions": {
"__start__": AnyVersion(),
@@ -2936,8 +2932,10 @@ async def test_send_dedupe_on_resume(
if checkpoint_during:
assert history == expected_history
else:
assert history[0] == expected_history[0]
assert history[1] == expected_history[2]
assert history[0] == expected_history[0]._replace(
parent_config=history[1].config
)
assert history[1] == expected_history[2]._replace(parent_config=None)
async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) -> None: