mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-18 14:07:31 +02:00
feat(checkpoint-postgres): add prune to PostgresSaver and AsyncPostgresSaver
`delete` removes every checkpoint, write and blob of the given threads. `keep_latest` keeps the latest checkpoint of each namespace, with its pending writes and the blobs it references, and deletes everything older, including blobs only those checkpoints used. Checkpoints saved while the prune runs are left alone. `keep_latest` refuses threads that use a DeltaChannel, since a DeltaChannel rebuilds the latest value from older checkpoints; nothing is pruned when any requested thread uses one.
This commit is contained in:
Elior Nataf Lackritz
@@ -27,10 +27,19 @@ 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_REFERENCED_BLOBS_SQL,
|
||||
PRUNE_STRATEGIES,
|
||||
PRUNE_WRITES_SQL,
|
||||
BasePostgresSaver,
|
||||
_build_delta_stage1_sql,
|
||||
_build_delta_stage2_sql,
|
||||
_DeltaStage2Row,
|
||||
_garbage_blobs,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
|
||||
@@ -401,6 +410,82 @@ 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 (plus its
|
||||
parent when it predates v4 and has pending sends), with its pending
|
||||
writes and the blobs it references, and deletes everything older. 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"` 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 while the
|
||||
batch runs is refused at its turn, 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`.
|
||||
"""
|
||||
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:
|
||||
self._refuse_delta_threads(cur, 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))
|
||||
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)}"
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
@@ -27,10 +27,19 @@ 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_REFERENCED_BLOBS_SQL,
|
||||
PRUNE_STRATEGIES,
|
||||
PRUNE_WRITES_SQL,
|
||||
BasePostgresSaver,
|
||||
_build_delta_stage1_sql,
|
||||
_build_delta_stage2_sql,
|
||||
_DeltaStage2Row,
|
||||
_garbage_blobs,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
|
||||
@@ -360,6 +369,71 @@ 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`.
|
||||
"""
|
||||
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 self._refuse_delta_threads(cur, 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))
|
||||
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)}"
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
@@ -688,6 +762,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]:
|
||||
|
||||
@@ -160,6 +160,114 @@ 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`), and keep_latest works one namespace
|
||||
# at a time.
|
||||
|
||||
# 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.
|
||||
PRUNE_DELTA_THREADS_SQL = """
|
||||
SELECT thread_id FROM checkpoints
|
||||
WHERE thread_id = ANY(%s) AND metadata ? 'counters_since_delta_snapshot'
|
||||
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 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.
|
||||
PRUNE_HEADS_SQL = f"""
|
||||
SELECT
|
||||
h.checkpoint_ns,
|
||||
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
|
||||
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 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_CHECKPOINTS_SQL = """
|
||||
WITH deleted AS (
|
||||
DELETE FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %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_WRITES_SQL = """
|
||||
DELETE FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id < %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)
|
||||
"""
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
class _DeltaStage2Row(TypedDict, total=False):
|
||||
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
|
||||
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
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,
|
||||
) -> 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=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 = "") -> 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)
|
||||
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]]
|
||||
) -> None:
|
||||
pending = [action]
|
||||
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:
|
||||
await self.cur.execute(query, params, **kwargs)
|
||||
if query is sql and pending:
|
||||
await pending.pop()()
|
||||
|
||||
@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]}},
|
||||
)
|
||||
else:
|
||||
first = await _put(
|
||||
saver, thread_id, None, 0, {"m": _DeltaSnapshot([1])}, {"m": "1"}
|
||||
)
|
||||
await _put(saver, thread_id, first, 1, {}, {"m": "2"})
|
||||
|
||||
|
||||
@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]
|
||||
assert await _rows(saver, other_id) == other_before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("evidence", ["counters", "snapshot"])
|
||||
@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"]
|
||||
|
||||
|
||||
@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"],
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,367 @@
|
||||
import operator
|
||||
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, 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"),
|
||||
]
|
||||
|
||||
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 _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))]
|
||||
|
||||
|
||||
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 _run_once_after(saver: PostgresSaver, sql: str, action: Callable[[], Any]) -> None:
|
||||
pending = [action]
|
||||
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:
|
||||
self.cur.execute(query, params, **kwargs)
|
||||
if query is sql and pending:
|
||||
pending.pop()()
|
||||
|
||||
@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")) == 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")) == 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")) == 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)
|
||||
namespaces = {
|
||||
c.config["configurable"]["checkpoint_ns"]
|
||||
for c in checkpointer.list(_config("pruned"))
|
||||
}
|
||||
|
||||
checkpointer.prune(["pruned"])
|
||||
|
||||
assert len(_checkpoint_ids(checkpointer, "pruned")) == len(namespaces)
|
||||
assert graph.get_state(_config("pruned"), subgraphs=True) == 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
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
@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")) == 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
|
||||
Reference in New Issue
Block a user