mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
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:
@@ -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",
|
||||
|
||||
@@ -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"],
|
||||
)
|
||||
@@ -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"]
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -122,7 +122,6 @@ RESERVED = {
|
||||
ERROR,
|
||||
NO_WRITES,
|
||||
SCHEDULED,
|
||||
TASKS,
|
||||
# reserved config.configurable keys
|
||||
CONFIG_KEY_SEND,
|
||||
CONFIG_KEY_READ,
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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", []),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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,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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user