WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-16 13:15:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(checkpoint-postgres): delete finished subgraph calls and prune each thread on one snapshot

Browse Source 
keep_latest now deletes a subgraph call's namespace together with the root
checkpoint the call started from. A checkpointer=True subgraph keeps its
latest checkpoint, since its one namespace holds state across calls, and a
call whose root checkpoint isn't saved yet is left alone.

Each thread is pruned in one REPEATABLE READ transaction that deletes
everything its snapshot sees except the kept checkpoints, so what a run
commits meanwhile survives whatever its ID. Writes go by the checkpoints
deleted, since a run can commit a superstep's writes before its checkpoint.

The per-thread statements live in one generator shared by the sync and
async savers, run unprepared: a cached generic plan checks an = ANY array
one element per row.
This commit is contained in:
Elior Nataf Lackritz
2026-10-14 13:40:06 -04:00
parent b46c897f4f
commit 389b3c1d1f
5 changed files with 324 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+28 -46
View File
@@ -27,20 +27,15 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE,
PRUNE_BLOBS_SQL,
PRUNE_CHECKPOINTS_SQL,
PRUNE_DELETE_THREADS_SQL,
PRUNE_DELTA_THREADS_SQL,
PRUNE_HEADS_SQL,
PRUNE_PARENTS_SQL,
PRUNE_REFERENCED_BLOBS_SQL,
PRUNE_STRATEGIES,
PRUNE_WRITES_SQL,
BasePostgresSaver,
_build_delta_stage1_sql,
_build_delta_stage2_sql,
_DeltaStage2Row,
_garbage_blobs,
_execute,
_keep_latest,
_refuse_delta_threads,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -422,13 +417,17 @@ class PostgresSaver(BasePostgresSaver):
`"delete"` removes every checkpoint, write and blob of the threads, in
one transaction.
`"keep_latest"` keeps the latest checkpoint of each namespace (plus its
parent when it predates v4 and has pending sends), with its pending
writes and the blobs it references, and deletes everything older. A
kept checkpoint whose parent it deletes is left with no parent. Each
thread is pruned in its own transaction. Checkpoints saved while it runs
are kept, so a run appending to the thread meanwhile is safe; forking or
replaying from a checkpoint being deleted is not.
`"keep_latest"` keeps the latest checkpoint of each namespace (the one
`get_tuple` loads, plus its parent when it predates v4 and has pending
sends), with its pending writes and the blobs it references, and
deletes the rest. A kept checkpoint whose parent it deletes is left with
no parent. A subgraph call's namespace is deleted with the root
checkpoint the call started from, except a `checkpointer=True`
subgraph's, which keeps its latest checkpoint. Each thread is pruned in
its own transaction, on one snapshot of it, so a run appending to the
thread meanwhile is safe; forking or replaying from a checkpoint being
deleted is not. A thread that fails because another transaction changed
it leaves the threads before it pruned, and retrying is safe.
`"keep_latest"` refuses threads that use a `DeltaChannel`, since it
rebuilds the latest value from older checkpoints: if any requested
@@ -444,6 +443,11 @@ class PostgresSaver(BasePostgresSaver):
ValueError: If `strategy` is unknown.
NotImplementedError: If `strategy` is `"keep_latest"` and a thread
uses a `DeltaChannel`.
psycopg.errors.SerializationFailure: If another transaction, such
as a second prune, changes a thread's rows while `"keep_latest"`
prunes it.
psycopg.errors.DeadlockDetected: If a second prune of a thread runs
at the same time.
"""
if isinstance(thread_ids, str):
raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
@@ -458,38 +462,16 @@ class PostgresSaver(BasePostgresSaver):
cur.execute(sql, (thread_ids,))
return
with self._cursor() as cur:
self._refuse_delta_threads(cur, thread_ids)
_execute(cur, _refuse_delta_threads(thread_ids))
for thread_id in thread_ids:
with self._cursor() as cur, cur.connection.transaction():
cur.execute(PRUNE_HEADS_SQL, (thread_id, thread_id, thread_id))
heads = cur.fetchall()
# A run may have written this thread's first DeltaChannel
# checkpoint since the batch check. Checking after reading the
# heads sees it before anything it rebuilds from is deleted.
self._refuse_delta_threads(cur, [thread_id])
for head in heads:
ns, keep_from = head["checkpoint_ns"], head["keep_from"]
cur.execute(PRUNE_CHECKPOINTS_SQL, (thread_id, ns, keep_from))
freed = cur.fetchall()
cur.execute(PRUNE_WRITES_SQL, (thread_id, ns, keep_from))
cur.execute(
PRUNE_PARENTS_SQL, (thread_id, ns, keep_from, keep_from)
)
if not freed:
continue
cur.execute(PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns))
garbage = _garbage_blobs(freed, cur.fetchall())
for channel, versions in garbage.items():
cur.execute(PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions))
def _refuse_delta_threads(
self, cur: Cursor[DictRow], thread_ids: Sequence[str]
) -> None:
cur.execute(PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids))
if delta := sorted(r["thread_id"] for r in cur.fetchall()):
raise NotImplementedError(
f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}"
)
with self._cursor() as cur:
if self.pipe:
# The transaction has to start with its isolation level, and
# in pipeline mode it would take in what ran since the last
# sync.
self.pipe.sync()
with cur.connection.transaction():
_execute(cur, _keep_latest(thread_id))
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+17 -41
View File
@@ -27,20 +27,15 @@ from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE,
PRUNE_BLOBS_SQL,
PRUNE_CHECKPOINTS_SQL,
PRUNE_DELETE_THREADS_SQL,
PRUNE_DELTA_THREADS_SQL,
PRUNE_HEADS_SQL,
PRUNE_PARENTS_SQL,
PRUNE_REFERENCED_BLOBS_SQL,
PRUNE_STRATEGIES,
PRUNE_WRITES_SQL,
BasePostgresSaver,
_aexecute,
_build_delta_stage1_sql,
_build_delta_stage2_sql,
_DeltaStage2Row,
_garbage_blobs,
_keep_latest,
_refuse_delta_threads,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -389,6 +384,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
ValueError: If `strategy` is unknown.
NotImplementedError: If `strategy` is `"keep_latest"` and a thread
uses a `DeltaChannel`.
psycopg.errors.SerializationFailure: If another transaction, such
as a second prune, changes a thread's rows while `"keep_latest"`
prunes it.
psycopg.errors.DeadlockDetected: If a second prune of a thread runs
at the same time.
"""
if isinstance(thread_ids, str):
raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
@@ -403,40 +403,16 @@ class AsyncPostgresSaver(BasePostgresSaver):
await cur.execute(sql, (thread_ids,))
return
async with self._cursor() as cur:
await self._refuse_delta_threads(cur, thread_ids)
await _aexecute(cur, _refuse_delta_threads(thread_ids))
for thread_id in thread_ids:
async with self._cursor() as cur, cur.connection.transaction():
await cur.execute(PRUNE_HEADS_SQL, (thread_id, thread_id, thread_id))
heads = await cur.fetchall()
# A run may have written this thread's first DeltaChannel
# checkpoint since the batch check. Checking after reading the
# heads sees it before anything it rebuilds from is deleted.
await self._refuse_delta_threads(cur, [thread_id])
for head in heads:
ns, keep_from = head["checkpoint_ns"], head["keep_from"]
await cur.execute(PRUNE_CHECKPOINTS_SQL, (thread_id, ns, keep_from))
freed = await cur.fetchall()
await cur.execute(PRUNE_WRITES_SQL, (thread_id, ns, keep_from))
await cur.execute(
PRUNE_PARENTS_SQL, (thread_id, ns, keep_from, keep_from)
)
if not freed:
continue
await cur.execute(PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns))
garbage = _garbage_blobs(freed, await cur.fetchall())
for channel, versions in garbage.items():
await cur.execute(
PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions)
)
async def _refuse_delta_threads(
self, cur: AsyncCursor[DictRow], thread_ids: Sequence[str]
) -> None:
await cur.execute(PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids))
if delta := sorted(r["thread_id"] for r in await cur.fetchall()):
raise NotImplementedError(
f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}"
)
async with self._cursor() as cur:
if self.pipe:
# The transaction has to start with its isolation level, and
# in pipeline mode it would take in what ran since the last
# sync.
await self.pipe.sync()
async with cur.connection.transaction():
await _aexecute(cur, _keep_latest(thread_id))
@asynccontextmanager
async def _cursor(
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+117 -21
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import random
import warnings
from collections.abc import Mapping, Sequence
from collections.abc import Generator, Mapping, Sequence
from importlib.metadata import version as get_version
from typing import Any, TypedDict, cast
@@ -17,6 +17,8 @@ from langgraph.checkpoint.base import (
writes_sort_key,
)
from langgraph.checkpoint.serde.types import TASKS
from psycopg import AsyncCursor, Cursor
from psycopg.rows import DictRow
from psycopg.types.json import Jsonb
# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal
@@ -171,8 +173,7 @@ PRUNE_DELETE_THREADS_SQL = (
# The prune statements below never join two sets of a thread's rows: a join is
# planned from per-thread row estimates, which are worst for the few huge
# threads worth pruning, and can go quadratic. They look rows up by primary key
# instead (`LATERAL ... LIMIT 1`, `= ANY`), and keep_latest works one namespace
# at a time.
# instead (`LATERAL ... LIMIT 1`, `= ANY`).
# A `DeltaChannel` leaves update counters in checkpoint metadata, or a
# snapshot: an inline `true` marker with a blob behind it, which a plain `True`
@@ -199,19 +200,26 @@ CROSS JOIN LATERAL (
WHERE c.thread_id = ANY(%s) AND v.value = 'true'::jsonb
"""
# Per namespace, the oldest checkpoint keep_latest keeps: the newest one, or
# its parent when the newest predates v4 and reads its pending sends from the
# parent's writes.
# Per namespace: the latest checkpoint, the parent it reads its pending sends
# from when it predates v4, and the root graph's checkpoint its subgraph call
# started from, with whether that one is saved.
PRUNE_HEADS_SQL = f"""
SELECT
h.checkpoint_ns,
h.checkpoint_id AS head,
CASE WHEN (c.checkpoint ->> 'v')::int < 4 AND EXISTS (
SELECT 1 FROM checkpoint_writes w
WHERE w.thread_id = %s
AND w.checkpoint_ns = h.checkpoint_ns
AND w.checkpoint_id = c.parent_checkpoint_id
AND w.channel = '{TASKS}'
) THEN c.parent_checkpoint_id ELSE h.checkpoint_id END AS keep_from
) THEN c.parent_checkpoint_id END AS sends_parent,
c.metadata -> 'parents' ->> '' AS root,
EXISTS (
SELECT 1 FROM checkpoints r
WHERE r.thread_id = %s AND r.checkpoint_ns = ''
AND r.checkpoint_id = c.metadata -> 'parents' ->> ''
) AS root_saved
FROM (
SELECT DISTINCT ON (checkpoint_ns) checkpoint_ns, checkpoint_id
FROM checkpoints
@@ -219,19 +227,31 @@ FROM (
ORDER BY checkpoint_ns DESC, checkpoint_id DESC
) AS h
CROSS JOIN LATERAL (
SELECT checkpoint, parent_checkpoint_id FROM checkpoints
SELECT checkpoint, parent_checkpoint_id, metadata FROM checkpoints
WHERE thread_id = %s
AND checkpoint_ns = h.checkpoint_ns
AND checkpoint_id = h.checkpoint_id
) AS c
"""
# Deleting only below the newest checkpoint read above leaves alone whatever a
# concurrent run saves meanwhile.
PRUNE_SNAPSHOT_SQL = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
PRUNE_UNKEPT_SQL = """
SELECT checkpoint_id FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
"""
# A run saves a superstep's writes in the background, where they can commit
# before their checkpoint does, so they go by the checkpoints deleted.
PRUNE_WRITES_SQL = """
DELETE FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = ANY(%s)
"""
PRUNE_CHECKPOINTS_SQL = """
WITH deleted AS (
DELETE FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %s
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
RETURNING checkpoint -> 'channel_versions' AS versions
)
SELECT DISTINCT v.key AS channel, v.value AS version
@@ -239,15 +259,10 @@ FROM deleted
CROSS JOIN LATERAL jsonb_each_text(deleted.versions) AS v
"""
PRUNE_WRITES_SQL = """
DELETE FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %s
"""
PRUNE_PARENTS_SQL = """
UPDATE checkpoints SET parent_checkpoint_id = NULL
WHERE thread_id = %s AND checkpoint_ns = %s
AND checkpoint_id >= %s AND parent_checkpoint_id < %s
AND checkpoint_id = ANY(%s) AND parent_checkpoint_id <> ALL(%s)
"""
PRUNE_REFERENCED_BLOBS_SQL = """
@@ -262,14 +277,17 @@ DELETE FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s AND version = ANY(%s)
"""
PRUNE_SUBGRAPH_CALLS_SQL = (
"DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
"DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
"DELETE FROM checkpoint_blobs WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
)
def _garbage_blobs(
freed: Sequence[Mapping[str, Any]], referenced: Sequence[Mapping[str, Any]]
) -> dict[str, list[str]]:
"""Only blobs a deleted checkpoint used: a put that commits after `referenced`
was read isn't in it, so deleting every blob missing from it would take that
put's new blobs.
"""
"""Blobs a deleted checkpoint used and no kept one does, by channel."""
kept = {(r["channel"], r["version"]) for r in referenced}
garbage: dict[str, list[str]] = {}
for r in freed:
@@ -278,6 +296,84 @@ def _garbage_blobs(
return garbage
_Statements = Generator[tuple[str, Sequence[Any], bool], list[DictRow], None]
def _refuse_delta_threads(thread_ids: Sequence[str]) -> _Statements:
rows = yield PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids), True
if delta := sorted(r["thread_id"] for r in rows):
raise NotImplementedError(
f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}"
)
def _keep_latest(thread_id: str) -> _Statements:
"""The statements of one thread's keep_latest transaction, for `_execute`
and `_aexecute` to run. Each comes with whether its rows are sent back.
They all read one snapshot, so what a run saves meanwhile is neither seen
nor deleted, whatever its checkpoint ID.
"""
yield PRUNE_SNAPSHOT_SQL, (), False
heads = yield PRUNE_HEADS_SQL, (thread_id,) * 4, True
# Catches a thread that started using a DeltaChannel after the batch check.
yield from _refuse_delta_threads([thread_id])
kept = {
h["checkpoint_ns"]: [i for i in (h["head"], h["sends_parent"]) if i]
for h in heads
}
# A subgraph call's namespace has its task ID after a ":", unlike a
# `checkpointer=True` subgraph's, whose state outlives each call. A call
# goes with the root checkpoint it started from, but only once that one is
# saved: runs save in the background, so a call can commit first.
calls = {
h["checkpoint_ns"]
for h in heads
if ":" in h["checkpoint_ns"]
and h["root_saved"]
and h["root"] not in kept.get("", ())
}
if calls:
for sql in PRUNE_SUBGRAPH_CALLS_SQL:
yield sql, (thread_id, list(calls)), False
for ns, ids in kept.items():
if ns in calls:
continue
unkept = yield PRUNE_UNKEPT_SQL, (thread_id, ns, ids), True
if unkept:
deleted = [r["checkpoint_id"] for r in unkept]
yield PRUNE_WRITES_SQL, (thread_id, ns, deleted), False
freed = yield PRUNE_CHECKPOINTS_SQL, (thread_id, ns, ids), True
yield PRUNE_PARENTS_SQL, (thread_id, ns, ids, ids), False
if freed:
referenced = yield PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns), True
for channel, versions in _garbage_blobs(freed, referenced).items():
yield PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions), False
# Unprepared, so Postgres plans each statement knowing its arrays: a cached
# generic plan checks a row against an `= ANY(...)` array one element at a time.
def _execute(cur: Cursor[DictRow], statements: _Statements) -> None:
try:
sql, params, fetch = next(statements)
while True:
cur.execute(sql, params, prepare=False)
sql, params, fetch = statements.send(cur.fetchall() if fetch else [])
except StopIteration:
pass
async def _aexecute(cur: AsyncCursor[DictRow], statements: _Statements) -> None:
try:
sql, params, fetch = next(statements)
while True:
await cur.execute(sql, params, prepare=False)
rows = await cur.fetchall() if fetch else []
sql, params, fetch = statements.send(rows)
except StopIteration:
pass
class _DeltaStage2Row(TypedDict, total=False):
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
libs/checkpoint-postgres/tests/test_prune.py
+110 -3
View File
@@ -74,6 +74,7 @@ async def _put(
ns: str = "",
metadata: dict[str, Any] | None = None,
v: int = 4,
checkpoint_id: str | None = None,
) -> dict:
config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ns}}
if parent is not None:
@@ -83,7 +84,7 @@ async def _put(
checkpoint = empty_checkpoint()
checkpoint.update(
v=v,
id=str(uuid6(clock_seq=step)),
id=checkpoint_id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=versions,
)
@@ -98,7 +99,12 @@ async def _put(
)
async def _three_checkpoints(saver: Any, thread_id: str, ns: str = "") -> dict:
async def _three_checkpoints(
saver: Any,
thread_id: str,
ns: str = "",
metadata: dict[str, Any] | None = None,
) -> dict:
parent = None
steps = [
({"a": [1], "s": ["shared"]}, {"a": "1", "s": "1"}),
@@ -106,7 +112,9 @@ async def _three_checkpoints(saver: Any, thread_id: str, ns: str = "") -> dict:
({"a": [1, 2, 3]}, {"a": "3", "s": "1"}),
]
for step, (values, versions) in enumerate(steps):
parent = await _put(saver, thread_id, parent, step, values, versions, ns=ns)
parent = await _put(
saver, thread_id, parent, step, values, versions, ns=ns, metadata=metadata
)
await _call(saver, "put_writes", parent, [("a", step)], str(uuid4()))
return parent
@@ -431,3 +439,102 @@ async def test_keep_latest_keeps_the_history_of_a_delta_checkpoint_saved_after_t
)
assert history["a"]["seed"] == [1, 2, 3]
assert [value for _, _, value in history["a"]["writes"]] == [2]
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_deletes_a_subgraph_call_with_the_checkpoint_it_started_from(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
older = await _put(saver, thread_id, None, 0, {}, {})
head = await _put(saver, thread_id, older, 1, {}, {})
unsaved = str(uuid6(clock_seq=2))
started_from = {
"finished:1": older["configurable"]["checkpoint_id"],
"interrupted:2": head["configurable"]["checkpoint_id"],
"running:3": unsaved,
"checkpointer_true": older["configurable"]["checkpoint_id"],
}
for ns, root in started_from.items():
await _three_checkpoints(
saver, thread_id, ns, metadata={"parents": {"": root}}
)
await _call(saver, "prune", [thread_id])
rows = await _rows(saver, thread_id)
assert {ns for ns, *_ in rows["checkpoints"]} == {
"",
"interrupted:2",
"running:3",
"checkpointer_true",
}
assert {ns for ns, *_ in rows["writes"]} == {
"interrupted:2",
"running:3",
"checkpointer_true",
}
assert {ns for ns, *_ in rows["blobs"]} == {
"interrupted:2",
"running:3",
"checkpointer_true",
}
async def test_keep_latest_keeps_a_checkpoint_saved_meanwhile_whatever_its_id() -> None:
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
lower_id = "0" + head["configurable"]["checkpoint_id"][1:]
saved: list[dict] = []
async def save_a_child_with_a_lower_id() -> None:
child = await _put(
other,
thread_id,
head,
3,
{"a": [1, 2, 3, 4]},
{"a": "4", "s": "1"},
checkpoint_id=lower_id,
)
await _call(other, "put_writes", child, [("a", 4)], str(uuid4()))
saved.append(child)
_run_once_after(saver, PRUNE_HEADS_SQL, save_a_child_with_a_lower_id)
await saver.aprune([thread_id])
loaded = await saver.aget_tuple(saved[0])
assert loaded is not None, "the checkpoint saved during the prune was deleted"
assert loaded.checkpoint["channel_values"] == {
"a": [1, 2, 3, 4],
"s": ["shared"],
}
assert [value for _, _, value in loaded.pending_writes or []] == [4]
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_keeps_writes_saved_before_their_checkpoint(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
next_id = str(uuid6(clock_seq=3))
next_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": "",
"checkpoint_id": next_id,
}
}
await _call(saver, "put_writes", next_config, [("a", 4)], str(uuid4()))
await _call(saver, "prune", [thread_id])
await _put(
saver, thread_id, head, 3, {}, {"a": "3", "s": "1"}, checkpoint_id=next_id
)
loaded = await _call(saver, "get_tuple", next_config)
assert [value for _, _, value in loaded.pending_writes] == [4]
libs/langgraph/tests/test_prune_postgres.py
+52 -5
View File
@@ -94,6 +94,22 @@ def _graph(
return builder.compile(checkpointer=checkpointer)
def _counting_subgraph(name: str, *, checkpointer: bool | None) -> Any:
class Calls(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
calls: Annotated[list[str], operator.add]
def count(state: Calls) -> dict:
n = len(state["messages"])
reply = AIMessage(f"{name} call {len(state['calls']) + 1}", id=f"{name}-{n}")
return {"calls": [name], "messages": [reply]}
builder = StateGraph(Calls)
builder.add_node("count", count)
builder.add_edge(START, "count")
return builder.compile(checkpointer=checkpointer)
def _delta_graph(checkpointer: Any) -> Any:
builder = StateGraph(DeltaState)
builder.add_node("reply", _reply(ask=False))
@@ -113,6 +129,13 @@ def _checkpoint_ids(checkpointer: Any, thread_id: str) -> list[str]:
return [c.checkpoint["id"] for c in checkpointer.list(_config(thread_id))]
def _namespaces(checkpointer: Any, thread_id: str) -> list[str]:
return sorted(
c.config["configurable"]["checkpoint_ns"]
for c in checkpointer.list(_config(thread_id))
)
async def _acheckpoint_ids(checkpointer: Any, thread_id: str) -> list[str]:
return [c.checkpoint["id"] async for c in checkpointer.alist(_config(thread_id))]
@@ -256,14 +279,11 @@ def test_keep_latest_keeps_an_interrupted_subgraph_resumable(
graph.invoke(Command(resume=f"yes {i}"), _config(thread_id))
graph.invoke(_say(2), _config(thread_id))
before = graph.get_state(_config("pruned"), subgraphs=True)
namespaces = {
c.config["configurable"]["checkpoint_ns"]
for c in checkpointer.list(_config("pruned"))
}
interrupted = before.tasks[0].state.config["configurable"]["checkpoint_ns"]
checkpointer.prune(["pruned"])
assert len(_checkpoint_ids(checkpointer, "pruned")) == len(namespaces)
assert _namespaces(checkpointer, "pruned") == ["", interrupted]
assert graph.get_state(_config("pruned"), subgraphs=True) == _parentless(before)
for thread_id in ("pruned", "twin"):
graph.invoke(Command(resume="yes 2"), _config(thread_id))
@@ -273,6 +293,33 @@ def test_keep_latest_keeps_an_interrupted_subgraph_resumable(
)
@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_a_checkpointer_true_subgraph_but_not_finished_calls(
make_checkpointer: Any, durability: str
) -> None:
with make_checkpointer() as checkpointer:
parent = StateGraph(State)
parent.add_node("memory", _counting_subgraph("memory", checkpointer=True))
parent.add_node("call", _counting_subgraph("call", checkpointer=None))
parent.add_edge(START, "memory")
parent.add_edge("memory", "call")
graph = parent.compile(checkpointer=checkpointer)
for i in range(3):
for thread_id in ("pruned", "twin"):
graph.invoke(_say(i), _config(thread_id), durability=durability)
checkpointer.prune(["pruned"])
assert _namespaces(checkpointer, "pruned") == ["", "memory"]
for thread_id in ("pruned", "twin"):
graph.invoke(_say(3), _config(thread_id), durability=durability)
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
def test_keep_latest_keeps_a_turn_saved_after_it_read_the_heads() -> None:
with (
_checkpointer_postgres() as checkpointer,