ryan
c329d15f7f
Every graph step stores a full snapshot of the conversation, nothing pruned old ones, and deleting a conversation left its thread behind (about 1.6 MB per checkpoint at 200 turns). After each turn, only the newest CHECKPOINT_KEEP_PER_THREAD (default 20) checkpoints of that thread are kept, plus their pending writes; state reads only use the newest, and pruning a 200-checkpoint thread takes about 2 ms. Deleting a conversation also deletes its thread, without failing the request if the checkpoint store errors. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
"""Keeps the LangGraph checkpoint DB from growing without bound.
|
|
|
|
Every `AsyncSqliteSaver` row stores a full snapshot of the graph's channel
|
|
values (the complete message list so far), not a delta — so an unpruned
|
|
thread's checkpoints grow roughly linearly with turn count.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
|
|
DEFAULT_KEEP_PER_THREAD = 20
|
|
|
|
|
|
def keep_per_thread() -> int:
|
|
raw = os.getenv("CHECKPOINT_KEEP_PER_THREAD", str(DEFAULT_KEEP_PER_THREAD))
|
|
try:
|
|
value = int(raw)
|
|
except ValueError:
|
|
value = -1
|
|
if value < 1:
|
|
logging.warning(
|
|
"CHECKPOINT_KEEP_PER_THREAD=%r is invalid; falling back to %d",
|
|
raw,
|
|
DEFAULT_KEEP_PER_THREAD,
|
|
)
|
|
return DEFAULT_KEEP_PER_THREAD
|
|
return value
|
|
|
|
|
|
async def prune_thread_checkpoints(
|
|
checkpointer, thread_id: str, keep: int | None = None
|
|
) -> None:
|
|
"""Keep only the newest `keep` checkpoints (and their writes) for a thread.
|
|
|
|
Safe to call after every turn: each checkpoint row is a self-contained
|
|
channel snapshot rather than a delta, and `aget_tuple`/`aget_state` only
|
|
ever read the single newest row for a (thread_id, checkpoint_ns) pair, so
|
|
dropping older rows never affects state reads. Reuses the saver's own
|
|
lock so this never races its aput/aget_tuple calls.
|
|
"""
|
|
keep = keep_per_thread() if keep is None else keep
|
|
thread_id = str(thread_id)
|
|
conn = checkpointer.conn
|
|
async with checkpointer.lock, conn.cursor() as cur:
|
|
await cur.execute(
|
|
"SELECT DISTINCT checkpoint_ns FROM checkpoints WHERE thread_id = ?",
|
|
(thread_id,),
|
|
)
|
|
namespaces = [row[0] for row in await cur.fetchall()]
|
|
for ns in namespaces:
|
|
await cur.execute(
|
|
"""
|
|
DELETE FROM checkpoints
|
|
WHERE thread_id = ? AND checkpoint_ns = ?
|
|
AND checkpoint_id NOT IN (
|
|
SELECT checkpoint_id FROM checkpoints
|
|
WHERE thread_id = ? AND checkpoint_ns = ?
|
|
ORDER BY checkpoint_id DESC
|
|
LIMIT ?
|
|
)
|
|
""",
|
|
(thread_id, ns, thread_id, ns, keep),
|
|
)
|
|
await cur.execute(
|
|
"""
|
|
DELETE FROM writes
|
|
WHERE thread_id = ? AND checkpoint_ns = ?
|
|
AND checkpoint_id NOT IN (
|
|
SELECT checkpoint_id FROM checkpoints
|
|
WHERE thread_id = ? AND checkpoint_ns = ?
|
|
)
|
|
""",
|
|
(thread_id, ns, thread_id, ns),
|
|
)
|
|
await conn.commit()
|
|
|
|
|
|
async def delete_thread_checkpoints(checkpointer, thread_id: str) -> None:
|
|
"""Delete a conversation's checkpoint thread; never raises.
|
|
|
|
Called from conversation delete, which must still succeed even if the
|
|
checkpoint store is unreachable or the thread is already gone.
|
|
"""
|
|
try:
|
|
await checkpointer.adelete_thread(str(thread_id))
|
|
except Exception:
|
|
logging.exception("Failed to delete checkpoint thread %s", thread_id)
|