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

feat(checkpoint-postgres): add prune to PostgresSaver and AsyncPostgresSaver (#9290)

Browse Source 
Closes #8531.

`PostgresSaver` and `AsyncPostgresSaver` inherited `prune` / `aprune`
from the base class, which raises `NotImplementedError`, so long-lived
threads had no supported way to drop old checkpoints. This adds both
strategies. Thanks to @wangjiawei-vegetable for proposing this scope in
#8531 and #9247.

- **`delete`** removes every checkpoint, write and blob of the threads,
in one transaction.
- **`keep_latest`** keeps the latest checkpoint of each namespace (the
one `get_tuple` loads), with its pending writes and the blobs it
references, and deletes the rest, including blobs only the deleted
checkpoints used. A kept checkpoint whose parent is deleted is left with
no parent. A subgraph call's namespace goes with the root checkpoint the
call started from; a `checkpointer=True` subgraph keeps its latest
checkpoint, since its state outlives each call.
- **Each thread is pruned on one snapshot.** Its transaction runs at
`REPEATABLE READ` and deletes everything that snapshot sees except the
kept checkpoints, so a run appending to the thread at the same time is
safe, whatever its checkpoint IDs.
- **`keep_latest` refuses threads that use a `DeltaChannel`**
(`NotImplementedError` naming them, nothing pruned). A `DeltaChannel`
rebuilds its latest value from older checkpoints, so deleting them would
empty it. The saver can't see channel types, so it detects one by what
only a `DeltaChannel` leaves: `counters_since_delta_snapshot` in
metadata, a snapshot (an inline `true` with a blob behind it), or
`delta_writes_versioned`, which LangGraph.js sets on threads its
`DeltaChannel` graphs start. The check runs for the whole batch first
and again in each thread's transaction, so a thread that starts using a
`DeltaChannel` mid-batch is refused at its turn.

Unlike the `prune` in #6883, `keep_latest` here deletes the blobs older
checkpoints no longer share with the kept ones (most of the storage on a
long thread), and guards `DeltaChannel` threads.

Measured locally, with or without planner statistics: a 6,000-checkpoint
thread prunes in 100 to 200 ms, and one with 1,500 finished subgraph
calls in about 110 ms.

## Known limitations

- Two prunes of the same thread at once can fail one of them with a
serialization failure or a deadlock; retrying is safe.
- Forking or replaying from a checkpoint while it's being pruned is not
safe.
- Threads written by older releases can escape the `DeltaChannel` check:
before #9142 (Python 1.2.13) and langchain-ai/langgraphjs#2973 (JS
1.5.0), some `update_state` paths saved checkpoints with neither
counters nor a snapshot.
- Postgres reuses the freed space but doesn't shrink the table files;
that takes `VACUUM FULL` or `pg_repack`.
- `DeltaChannel`-aware pruning (keeping ancestors back to the last
snapshot) is a follow-up.

## Test plan

- [x] `libs/checkpoint-postgres/tests/test_prune.py`: 71 tests over the
sync and async savers (single connection, pool, pipeline), including the
conformance `prune` capability. Passes on Postgres 15 and 16.
- [x] `libs/langgraph/tests/test_prune_postgres.py`: 39 graph tests. A
pruned thread loads the same state and keeps running like an unpruned
twin (sync, async and exit durability, interrupts, a failed parallel
step, interrupted and `checkpointer=True` subgraphs), plus deterministic
races: a turn saved mid-prune, and threads switching to a `DeltaChannel`
before and after their own check.
- [x] Each guard (the snapshot, the writes and blob scopes, the subgraph
call rules, the parent reset, the `DeltaChannel` checks, the
single-string check, the pre-v4 pending sends, the pipeline sync) broken
one at a time: every break fails at least one test.
- [x] A thread pruned in a loop while it runs, in each durability mode,
ends identical to its twin; threads written by LangGraph.js 1.5.0 and
pruned from Python behave like their twins.
This commit is contained in:
Elior Nataf Lackritz
2026-10-14 15:09:51 -04:00
committed by GitHub
parent e5b422455d
commit 36bf3887f0
5 changed files with 1375 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+72
View File
@@ -27,10 +27,15 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE,
PRUNE_DELETE_THREADS_SQL,
PRUNE_STRATEGIES,
BasePostgresSaver,
_build_delta_stage1_sql,
_build_delta_stage2_sql,
_DeltaStage2Row,
_execute,
_keep_latest,
_refuse_delta_threads,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -401,6 +406,73 @@ class PostgresSaver(BasePostgresSaver):
(str(thread_id),),
)
def prune(
self,
thread_ids: Sequence[str],
*,
strategy: str = "keep_latest",
) -> None:
"""Prune checkpoints for the given threads.
`"delete"` removes every checkpoint, write and blob of the threads, in
one transaction.
`"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
thread does, nothing is pruned. A thread that starts using one before
its turn in the batch is refused then, after the threads before it.
Args:
thread_ids: The thread IDs to prune.
strategy: `"keep_latest"` or `"delete"`.
Raises:
TypeError: If `thread_ids` is a single string.
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")
if strategy not in PRUNE_STRATEGIES:
raise ValueError(f"Unknown prune strategy: {strategy!r}")
thread_ids = [str(t) for t in thread_ids]
if not thread_ids:
return
if strategy == "delete":
with self._cursor() as cur, cur.connection.transaction():
for sql in PRUNE_DELETE_THREADS_SQL:
cur.execute(sql, (thread_ids,))
return
with self._cursor() as cur:
_execute(cur, _refuse_delta_threads(thread_ids))
for thread_id in thread_ids:
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]]:
"""Create a database cursor as a context manager.
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+81
View File
@@ -27,10 +27,15 @@ from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE,
PRUNE_DELETE_THREADS_SQL,
PRUNE_STRATEGIES,
BasePostgresSaver,
_aexecute,
_build_delta_stage1_sql,
_build_delta_stage2_sql,
_DeltaStage2Row,
_keep_latest,
_refuse_delta_threads,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -360,6 +365,55 @@ class AsyncPostgresSaver(BasePostgresSaver):
(str(thread_id),),
)
async def aprune(
self,
thread_ids: Sequence[str],
*,
strategy: str = "keep_latest",
) -> None:
"""Prune checkpoints for the given threads.
See `PostgresSaver.prune` for what each strategy keeps.
Args:
thread_ids: The thread IDs to prune.
strategy: `"keep_latest"` or `"delete"`.
Raises:
TypeError: If `thread_ids` is a single string.
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")
if strategy not in PRUNE_STRATEGIES:
raise ValueError(f"Unknown prune strategy: {strategy!r}")
thread_ids = [str(t) for t in thread_ids]
if not thread_ids:
return
if strategy == "delete":
async with self._cursor() as cur, cur.connection.transaction():
for sql in PRUNE_DELETE_THREADS_SQL:
await cur.execute(sql, (thread_ids,))
return
async with self._cursor() as cur:
await _aexecute(cur, _refuse_delta_threads(thread_ids))
for thread_id in thread_ids:
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(
self, *, pipeline: bool = False
@@ -688,6 +742,33 @@ class AsyncPostgresSaver(BasePostgresSaver):
self.adelete_thread(thread_id), self.loop
).result()
def prune(
self,
thread_ids: Sequence[str],
*,
strategy: str = "keep_latest",
) -> None:
"""Prune checkpoints for the given threads.
See `PostgresSaver.prune` for what each strategy keeps.
Args:
thread_ids: The thread IDs to prune.
strategy: `"keep_latest"` or `"delete"`.
"""
try:
if asyncio.get_running_loop() is self.loop:
raise asyncio.InvalidStateError(
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
"different thread. From the main thread, use the async interface. "
"For example, use `await checkpointer.aprune(...)`."
)
except RuntimeError:
pass
return asyncio.run_coroutine_threadsafe(
self.aprune(thread_ids, strategy=strategy), self.loop
).result()
def get_delta_channel_history(
self, *, config: RunnableConfig, channels: Sequence[str]
) -> Mapping[str, DeltaChannelHistory]:
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+215 -1
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
@@ -160,6 +162,218 @@ INSERT_CHECKPOINT_WRITES_SQL = """
"""
PRUNE_STRATEGIES = ("keep_latest", "delete")
PRUNE_DELETE_THREADS_SQL = (
"DELETE FROM checkpoints WHERE thread_id = ANY(%s)",
"DELETE FROM checkpoint_blobs WHERE thread_id = ANY(%s)",
"DELETE FROM checkpoint_writes WHERE thread_id = ANY(%s)",
)
# 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`).
# A `DeltaChannel` leaves update counters in checkpoint metadata, or a
# snapshot: an inline `true` marker with a blob behind it, which a plain `True`
# value never has. LangGraph.js inlines no values, but marks the threads its
# DeltaChannel graphs start with `delta_writes_versioned`.
PRUNE_DELTA_THREADS_SQL = """
SELECT thread_id FROM checkpoints
WHERE thread_id = ANY(%s) AND (
metadata ? 'counters_since_delta_snapshot'
OR metadata ? 'delta_writes_versioned'
)
UNION
SELECT c.thread_id
FROM checkpoints c
CROSS JOIN LATERAL jsonb_each(c.checkpoint -> 'channel_values') AS v
CROSS JOIN LATERAL (
SELECT 1 FROM checkpoint_blobs b
WHERE b.thread_id = c.thread_id
AND b.checkpoint_ns = c.checkpoint_ns
AND b.channel = v.key
AND b.version = c.checkpoint -> 'channel_versions' ->> v.key
LIMIT 1
) AS snapshot
WHERE c.thread_id = ANY(%s) AND v.value = 'true'::jsonb
"""
# 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 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
WHERE thread_id = %s
ORDER BY checkpoint_ns DESC, checkpoint_id DESC
) AS h
CROSS JOIN LATERAL (
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
"""
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 <> ALL(%s)
RETURNING checkpoint -> 'channel_versions' AS versions
)
SELECT DISTINCT v.key AS channel, v.value AS version
FROM deleted
CROSS JOIN LATERAL jsonb_each_text(deleted.versions) AS v
"""
PRUNE_PARENTS_SQL = """
UPDATE checkpoints SET parent_checkpoint_id = NULL
WHERE thread_id = %s AND checkpoint_ns = %s
AND checkpoint_id = ANY(%s) AND parent_checkpoint_id <> ALL(%s)
"""
PRUNE_REFERENCED_BLOBS_SQL = """
SELECT DISTINCT v.key AS channel, v.value AS version
FROM checkpoints c
CROSS JOIN LATERAL jsonb_each_text(c.checkpoint -> 'channel_versions') AS v
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
"""
PRUNE_BLOBS_SQL = """
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]]:
"""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:
if (r["channel"], r["version"]) not in kept:
garbage.setdefault(r["channel"], []).append(r["version"])
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
+540
View File
@@ -0,0 +1,540 @@
import asyncio
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager
from typing import Any
from uuid import uuid4
import pytest
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.conformance import validate
from langgraph.checkpoint.conformance.initializer import checkpointer_test
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
from psycopg import AsyncConnection
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.postgres.base import (
PRUNE_DELTA_THREADS_SQL,
PRUNE_HEADS_SQL,
PRUNE_REFERENCED_BLOBS_SQL,
)
from tests.conftest import DEFAULT_POSTGRES_URI, DEFAULT_URI
from tests.test_async import _saver as _async_saver
from tests.test_sync import _saver as _sync_saver
SAVERS = ["base", "pool", "pipe", "aio_base", "aio_pool", "aio_pipe"]
@asynccontextmanager
async def _saver(name: str):
if name.startswith("aio_"):
async with _async_saver(name.removeprefix("aio_")) as saver:
yield saver
else:
with _sync_saver(name) as saver:
yield saver
async def _call(saver: Any, method: str, *args: Any, **kwargs: Any) -> Any:
if isinstance(saver, AsyncPostgresSaver):
return await getattr(saver, f"a{method}")(*args, **kwargs)
return getattr(saver, method)(*args, **kwargs)
async def _rows(saver: Any, thread_id: str) -> dict[str, list[tuple]]:
queries = {
"checkpoints": "SELECT checkpoint_ns, checkpoint_id FROM checkpoints"
" WHERE thread_id = %s ORDER BY 1, 2",
"writes": "SELECT checkpoint_ns, checkpoint_id, task_id, idx"
" FROM checkpoint_writes WHERE thread_id = %s ORDER BY 1, 2, 3, 4",
"blobs": "SELECT checkpoint_ns, channel, version FROM checkpoint_blobs"
" WHERE thread_id = %s ORDER BY 1, 2, 3",
}
rows = {}
for table, sql in queries.items():
if isinstance(saver, AsyncPostgresSaver):
async with saver._cursor() as cur:
await cur.execute(sql, (thread_id,))
rows[table] = [tuple(r.values()) for r in await cur.fetchall()]
else:
with saver._cursor() as cur:
cur.execute(sql, (thread_id,))
rows[table] = [tuple(r.values()) for r in cur.fetchall()]
return rows
async def _put(
saver: Any,
thread_id: str,
parent: dict | None,
step: int,
values: dict[str, Any],
versions: dict[str, str],
*,
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:
config["configurable"]["checkpoint_id"] = parent["configurable"][
"checkpoint_id"
]
checkpoint = empty_checkpoint()
checkpoint.update(
v=v,
id=checkpoint_id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=versions,
)
new_versions = {k: versions[k] for k in values}
return await _call(
saver,
"put",
config,
checkpoint,
{"source": "loop", "step": step, "parents": {}, **(metadata or {})},
new_versions,
)
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"}),
({"a": [1, 2]}, {"a": "2", "s": "1"}),
({"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, metadata=metadata
)
await _call(saver, "put_writes", parent, [("a", step)], str(uuid4()))
return parent
@asynccontextmanager
async def _second_saver(saver: AsyncPostgresSaver):
assert isinstance(saver.conn, AsyncConnection)
uri = DEFAULT_POSTGRES_URI + saver.conn.info.dbname
async with AsyncPostgresSaver.from_conn_string(uri) as other:
yield other
def _run_once_after(
saver: AsyncPostgresSaver,
sql: str,
action: Callable[[], Awaitable[Any]],
*,
nth: int = 1,
) -> None:
seen = 0
cursor = saver._cursor
class Hooked:
def __init__(self, cur: Any) -> None:
self.cur = cur
def __getattr__(self, name: str) -> Any:
return getattr(self.cur, name)
async def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
nonlocal seen
await self.cur.execute(query, params, **kwargs)
if query is sql:
seen += 1
if seen == nth:
await action()
@asynccontextmanager
async def hooked(**kwargs: Any):
async with cursor(**kwargs) as cur:
yield Hooked(cur)
saver._cursor = hooked
async def _delta_thread(saver: Any, thread_id: str, evidence: str) -> None:
if evidence == "counters":
first = await _put(saver, thread_id, None, 0, {"m": [1]}, {"m": "1"})
await _put(
saver,
thread_id,
first,
1,
{},
{"m": "2"},
metadata={"counters_since_delta_snapshot": {"m": [1, 1]}},
)
elif evidence == "snapshot":
first = await _put(
saver, thread_id, None, 0, {"m": _DeltaSnapshot([1])}, {"m": "1"}
)
await _put(saver, thread_id, first, 1, {}, {"m": "2"})
else:
marker = {"delta_writes_versioned": True}
first = await _put(
saver, thread_id, None, 0, {"m": [1]}, {"m": "1"}, metadata=marker
)
await _put(
saver, thread_id, first, 1, {"m": [1, 2]}, {"m": "2"}, metadata=marker
)
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_keeps_each_namespace_head_with_its_writes_and_blobs(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id, other_id = str(uuid4()), str(uuid4())
heads = {
ns: await _three_checkpoints(saver, thread_id, ns) for ns in ("", "sub:1")
}
await _three_checkpoints(saver, other_id)
before = {ns: await _call(saver, "get_tuple", h) for ns, h in heads.items()}
other_before = await _rows(saver, other_id)
await _call(saver, "prune", [thread_id])
rows = await _rows(saver, thread_id)
assert rows["checkpoints"] == [
(ns, heads[ns]["configurable"]["checkpoint_id"]) for ns in ("", "sub:1")
]
assert [(ns, cid) for ns, cid, *_ in rows["writes"]] == rows["checkpoints"]
assert rows["blobs"] == [
(ns, ch, ver)
for ns in ("", "sub:1")
for ch, ver in (("a", "3"), ("s", "1"))
]
for ns, head in heads.items():
assert await _call(saver, "get_tuple", head) == before[ns]._replace(
parent_config=None
)
assert await _rows(saver, other_id) == other_before
@pytest.mark.parametrize("evidence", ["counters", "snapshot", "versioned"])
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
name: str, evidence: str
) -> None:
async with _saver(name) as saver:
ordinary_id, delta_id = str(uuid4()), str(uuid4())
await _three_checkpoints(saver, ordinary_id)
await _delta_thread(saver, delta_id, evidence)
before = [await _rows(saver, t) for t in (ordinary_id, delta_id)]
with pytest.raises(NotImplementedError, match=delta_id):
await _call(saver, "prune", [ordinary_id, delta_id])
assert [await _rows(saver, t) for t in (ordinary_id, delta_id)] == before
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_prunes_a_thread_storing_a_plain_true(name: str) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
first = await _put(
saver, thread_id, None, 0, {"flag": True, "a": [1]}, {"flag": "1", "a": "1"}
)
head = await _put(
saver, thread_id, first, 1, {"flag": True}, {"flag": "2", "a": "1"}
)
await _call(saver, "prune", [thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", head["configurable"]["checkpoint_id"])
]
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_keeps_the_parent_holding_a_pre_v4_heads_sends(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
first = await _put(saver, thread_id, None, 0, {}, {}, v=1)
await _call(saver, "put_writes", first, [(TASKS, "send-1")], str(uuid4()))
second = await _put(saver, thread_id, first, 1, {}, {}, v=1)
head = await _put(saver, thread_id, second, 2, {}, {}, v=1)
await _call(saver, "put_writes", second, [(TASKS, "send-2")], str(uuid4()))
await _call(saver, "prune", [thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", second["configurable"]["checkpoint_id"]),
("", head["configurable"]["checkpoint_id"]),
]
loaded = await _call(saver, "get_tuple", head)
assert loaded.checkpoint["channel_values"][TASKS] == ["send-2"]
assert loaded.parent_config == second
assert (await _call(saver, "get_tuple", second)).parent_config is None
@pytest.mark.parametrize("name", SAVERS)
async def test_delete_removes_every_row_of_the_threads(name: str) -> None:
async with _saver(name) as saver:
ordinary_id, delta_id, other_id = str(uuid4()), str(uuid4()), str(uuid4())
await _three_checkpoints(saver, ordinary_id)
await _delta_thread(saver, delta_id, "counters")
await _three_checkpoints(saver, other_id)
other_before = await _rows(saver, other_id)
await _call(saver, "prune", [ordinary_id, delta_id], strategy="delete")
empty = {"checkpoints": [], "writes": [], "blobs": []}
assert await _rows(saver, ordinary_id) == empty
assert await _rows(saver, delta_id) == empty
assert await _rows(saver, other_id) == other_before
@pytest.mark.parametrize("name", SAVERS)
async def test_unknown_strategy_changes_nothing(name: str) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
await _three_checkpoints(saver, thread_id)
before = await _rows(saver, thread_id)
with pytest.raises(ValueError, match="keep_all"):
await _call(saver, "prune", [thread_id], strategy="keep_all")
assert await _rows(saver, thread_id) == before
async def test_async_saver_prune_runs_from_another_thread_only() -> None:
async with _saver("aio_base") as saver:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
with pytest.raises(asyncio.InvalidStateError):
saver.prune([thread_id])
await asyncio.to_thread(saver.prune, [thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", head["configurable"]["checkpoint_id"])
]
async def test_prune_conformance() -> None:
@checkpointer_test(name="AsyncPostgresSaver")
async def postgres_saver():
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
yield saver
report = await validate(postgres_saver, capabilities={"prune"})
prune = report.results["prune"]
if not prune.passed:
details = "\n".join(prune.failures)
pytest.fail(f"Capability prune failed (passed={prune.passed}):\n{details}")
@pytest.mark.parametrize("strategy", ["keep_latest", "delete"])
@pytest.mark.parametrize("name", ["base", "aio_base"])
async def test_prune_rejects_a_single_thread_id_string(
name: str, strategy: str
) -> None:
async with _saver(name) as saver:
for thread_id in ("ab", "a", "b"):
await _three_checkpoints(saver, thread_id)
before = [await _rows(saver, t) for t in ("ab", "a", "b")]
with pytest.raises(TypeError, match="not a str"):
await _call(saver, "prune", "ab", strategy=strategy)
assert [await _rows(saver, t) for t in ("ab", "a", "b")] == before
@pytest.mark.parametrize(
"after", [PRUNE_HEADS_SQL, PRUNE_REFERENCED_BLOBS_SQL], ids=["heads", "blobs"]
)
async def test_keep_latest_keeps_a_checkpoint_saved_while_it_runs(after: str) -> None:
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
saved: list[dict] = []
async def save_a_child() -> None:
saved.append(
await _put(
other, thread_id, head, 3, {"a": [1, 2, 3, 4]}, {"a": "4", "s": "1"}
)
)
_run_once_after(saver, after, save_a_child)
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 loaded.parent_config == head
async def test_keep_latest_refuses_a_thread_that_turns_delta_after_the_batch_check() -> (
None
):
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
async def save_a_delta_checkpoint() -> None:
await _put(
other,
thread_id,
head,
3,
{},
{"a": "4", "s": "1"},
metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
)
_run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint)
with pytest.raises(NotImplementedError, match=thread_id):
await saver.aprune([thread_id])
assert len((await _rows(saver, thread_id))["checkpoints"]) == 4
async def test_keep_latest_keeps_the_history_of_a_delta_checkpoint_saved_after_the_thread_check() -> (
None
):
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
saved: list[dict] = []
async def save_a_delta_checkpoint() -> None:
saved.append(
await _put(
other,
thread_id,
head,
3,
{},
{"a": "4", "s": "1"},
metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
)
)
_run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint, nth=2)
await saver.aprune([thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", head["configurable"]["checkpoint_id"]),
("", saved[0]["configurable"]["checkpoint_id"]),
]
history = await saver.aget_delta_channel_history(
config=saved[0], channels=["a"]
)
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
+467
View File
@@ -0,0 +1,467 @@
import operator
import sys
from collections.abc import Callable
from contextlib import contextmanager
from typing import Annotated, Any
import pytest
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres.base import PRUNE_DELTA_THREADS_SQL, PRUNE_HEADS_SQL
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer, add_messages
from langgraph.types import Command, StateSnapshot, interrupt
from tests.conftest import NO_DOCKER
from tests.conftest_checkpointer import (
DEFAULT_POSTGRES_URI,
_checkpointer_postgres,
_checkpointer_postgres_aio,
_checkpointer_postgres_aio_pipe,
_checkpointer_postgres_aio_pool,
_checkpointer_postgres_pipe,
_checkpointer_postgres_pool,
)
pytestmark = [
pytest.mark.anyio,
pytest.mark.skipif(NO_DOCKER, reason="needs Postgres"),
]
NEEDS_CONTEXTVARS = pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
SYNC = [
_checkpointer_postgres,
_checkpointer_postgres_pipe,
_checkpointer_postgres_pool,
]
ASYNC = [
_checkpointer_postgres_aio,
_checkpointer_postgres_aio_pipe,
_checkpointer_postgres_aio_pool,
]
class State(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
notes: Annotated[list[str], operator.add]
class DeltaState(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
def _reply(ask: bool):
def reply(state: dict) -> dict:
n = len(state["messages"])
answer = interrupt("approve?") if ask else "ok"
return {"messages": [AIMessage(f"{answer} {n}", id=f"reply-{n}")]}
return reply
def _graph(
checkpointer: Any,
*,
ask: bool = False,
fail: set[str] | None = None,
calls: dict[str, int] | None = None,
) -> Any:
def note(name: str):
def node(state: State) -> dict:
if calls is not None:
calls[name] = calls.get(name, 0) + 1
if fail is not None and name in fail:
fail.discard(name)
raise ValueError(f"{name} failed")
return {"notes": [name]}
return node
builder = StateGraph(State)
builder.add_node("left", note("left"))
builder.add_node("right", note("right"))
builder.add_node("reply", _reply(ask))
builder.add_edge(START, "left")
builder.add_edge(START, "right")
builder.add_edge(["left", "right"], "reply")
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))
builder.add_edge(START, "reply")
return builder.compile(checkpointer=checkpointer)
def _config(thread_id: str) -> RunnableConfig:
return {"configurable": {"thread_id": thread_id}}
def _say(i: int) -> dict:
return {"messages": [HumanMessage(f"hi {i}", id=f"hi-{i}")]}
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))]
def _parentless(snapshot: StateSnapshot) -> StateSnapshot:
return snapshot._replace(
parent_config=None,
tasks=tuple(
t._replace(state=_parentless(t.state))
if isinstance(t.state, StateSnapshot)
else t
for t in snapshot.tasks
),
)
def _run_once_after(
saver: PostgresSaver, sql: str, action: Callable[[], Any], *, nth: int = 1
) -> None:
seen = 0
cursor = saver._cursor
class Hooked:
def __init__(self, cur: Any) -> None:
self.cur = cur
def __getattr__(self, name: str) -> Any:
return getattr(self.cur, name)
def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
nonlocal seen
self.cur.execute(query, params, **kwargs)
if query is sql:
seen += 1
if seen == nth:
action()
@contextmanager
def hooked(**kwargs: Any):
with cursor(**kwargs) as cur:
yield Hooked(cur)
saver._cursor = hooked
@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_the_state_and_the_thread_runs_on_like_its_twin(
make_checkpointer: Any, durability: str
) -> None:
with make_checkpointer() as checkpointer:
graph = _graph(checkpointer)
for i in range(3):
for thread_id in ("pruned", "twin"):
graph.invoke(_say(i), _config(thread_id), durability=durability)
before = graph.get_state(_config("pruned"))
checkpointer.prune(["pruned"])
assert _checkpoint_ids(checkpointer, "pruned") == [
before.config["configurable"]["checkpoint_id"]
]
assert graph.get_state(_config("pruned")) == _parentless(before)
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
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_a_pending_interrupt_resumable(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
graph = _graph(checkpointer, ask=True)
for thread_id in ("pruned", "twin"):
for i in range(2):
graph.invoke(_say(i), _config(thread_id))
graph.invoke(Command(resume=f"yes {i}"), _config(thread_id))
graph.invoke(_say(2), _config(thread_id))
before = graph.get_state(_config("pruned"))
assert before.interrupts
checkpointer.prune(["pruned"])
assert graph.get_state(_config("pruned")) == _parentless(before)
for thread_id in ("pruned", "twin"):
graph.invoke(Command(resume="yes 2"), _config(thread_id))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_writes_of_tasks_that_finished_before_a_failure(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
calls: dict[str, int] = {}
fail: set[str] = set()
graph = _graph(checkpointer, fail=fail, calls=calls)
for thread_id in ("pruned", "twin"):
for i in range(2):
graph.invoke(_say(i), _config(thread_id))
fail.add("right")
with pytest.raises(ValueError, match="right failed"):
graph.invoke(_say(2), _config(thread_id))
before = graph.get_state(_config("pruned"))
checkpointer.prune(["pruned"])
assert graph.get_state(_config("pruned")) == _parentless(before)
left_runs = calls["left"]
graph.invoke(None, _config("pruned"))
assert calls["left"] == left_runs, "a task that had finished ran again"
graph.invoke(None, _config("twin"))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_an_interrupted_subgraph_resumable(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
sub = StateGraph(State)
sub.add_node("reply", _reply(ask=True))
sub.add_edge(START, "reply")
parent = StateGraph(State)
parent.add_node("sub", sub.compile())
parent.add_edge(START, "sub")
graph = parent.compile(checkpointer=checkpointer)
for thread_id in ("pruned", "twin"):
for i in range(2):
graph.invoke(_say(i), _config(thread_id))
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)
interrupted = before.tasks[0].state.config["configurable"]["checkpoint_ns"]
checkpointer.prune(["pruned"])
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))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@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,
PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
) as pruner,
):
graph = _graph(checkpointer)
for i in range(3):
for thread_id in ("pruned", "twin"):
graph.invoke(_say(i), _config(thread_id))
_run_once_after(
pruner, PRUNE_HEADS_SQL, lambda: graph.invoke(_say(3), _config("pruned"))
)
pruner.prune(["pruned"])
graph.invoke(_say(3), _config("twin"))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
def test_keep_latest_refuses_a_thread_that_switches_to_a_delta_channel_mid_batch() -> (
None
):
with (
_checkpointer_postgres() as checkpointer,
PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
) as pruner,
):
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
for thread_id in ("switched", "twin"):
graph.invoke(_say(i), _config(thread_id))
_run_once_after(
pruner,
PRUNE_DELTA_THREADS_SQL,
lambda: delta_graph.invoke(_say(3), _config("switched")),
)
with pytest.raises(NotImplementedError, match="switched"):
pruner.prune(["switched"])
delta_graph.invoke(_say(3), _config("twin"))
assert (
delta_graph.get_state(_config("switched")).values
== delta_graph.get_state(_config("twin")).values
)
def test_keep_latest_keeps_a_delta_channel_turn_saved_after_the_thread_check() -> None:
with (
_checkpointer_postgres() as checkpointer,
PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
) as pruner,
):
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
for thread_id in ("switched", "twin"):
graph.invoke(_say(i), _config(thread_id))
head = _checkpoint_ids(checkpointer, "switched")[0]
_run_once_after(
pruner,
PRUNE_DELTA_THREADS_SQL,
lambda: delta_graph.invoke(_say(3), _config("switched")),
nth=2,
)
pruner.prune(["switched"])
assert _checkpoint_ids(checkpointer, "switched")[-1] == head
delta_graph.invoke(_say(3), _config("twin"))
assert (
delta_graph.get_state(_config("switched")).values
== delta_graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
graph.invoke(_say(i), _config("ordinary"))
delta_graph.invoke(_say(i), _config("delta"))
before = {t: _checkpoint_ids(checkpointer, t) for t in ("ordinary", "delta")}
delta_before = delta_graph.get_state(_config("delta"))
with pytest.raises(NotImplementedError, match="delta"):
checkpointer.prune(["ordinary", "delta"])
assert {t: _checkpoint_ids(checkpointer, t) for t in before} == before
assert delta_graph.get_state(_config("delta")) == delta_before
@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("make_checkpointer", ASYNC)
async def test_akeep_latest_keeps_a_pending_interrupt_resumable(
make_checkpointer: Any,
) -> None:
async with make_checkpointer() as checkpointer:
graph = _graph(checkpointer, ask=True)
for thread_id in ("pruned", "twin"):
for i in range(2):
await graph.ainvoke(_say(i), _config(thread_id))
await graph.ainvoke(Command(resume=f"yes {i}"), _config(thread_id))
await graph.ainvoke(_say(2), _config(thread_id))
before = await graph.aget_state(_config("pruned"))
await checkpointer.aprune(["pruned"])
assert await _acheckpoint_ids(checkpointer, "pruned") == [
before.config["configurable"]["checkpoint_id"]
]
assert await graph.aget_state(_config("pruned")) == _parentless(before)
for thread_id in ("pruned", "twin"):
await graph.ainvoke(Command(resume="yes 2"), _config(thread_id))
assert (await graph.aget_state(_config("pruned"))).values == (
await graph.aget_state(_config("twin"))
).values
@pytest.mark.parametrize("make_checkpointer", ASYNC)
async def test_akeep_latest_refuses_a_batch_with_a_delta_channel_thread(
make_checkpointer: Any,
) -> None:
async with make_checkpointer() as checkpointer:
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
await graph.ainvoke(_say(i), _config("ordinary"))
await delta_graph.ainvoke(_say(i), _config("delta"))
before = {
t: await _acheckpoint_ids(checkpointer, t) for t in ("ordinary", "delta")
}
with pytest.raises(NotImplementedError, match="delta"):
await checkpointer.aprune(["ordinary", "delta"])
assert {t: await _acheckpoint_ids(checkpointer, t) for t in before} == before