mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-14 12:15:25 +02:00
2b93e8b79f
Fixes langchain-ai/langgraph#8382 `DeltaChannel` rebuilds its value by replaying ancestor writes through the reducer. Every saver ordered a checkpoint's writes by `(task_id, idx)`, but live execution applies a super-step's writes in task-path order (`apply_writes` sorts tasks by `task_path_str(task.path[:3])`). `task_id` is a hash of the path, so when two or more tasks write one `DeltaChannel` in the same super-step, replay applies them in an arbitrary permutation. Reducers must be batching-invariant, not order-invariant, so the permutation changes the value: `get_state` disagrees with what `invoke` returned, and continuing the thread saves the reordered value as the base for every later write. ## Fix Replay orders by `(task_path, task_id, idx)`, the order `SELECT_PENDING_SENDS_SQL` already uses for sends. `apply_writes` and every saver now sort with one `writes_sort_key` from `langgraph.checkpoint.base`, in Python, so live and replay can't drift apart and don't depend on a database collation. `apply_writes` sorts by the full task path, the string `put_writes` stores, instead of `path[:3]`. The paths that reach `apply_writes` no longer carry task ids, so the live order doesn't change. `BaseCheckpointSaver`'s default `get_delta_channel_history` can't sort by path, since `PendingWrite` has no task path, so it replays each checkpoint's writes in `get_tuple` order. `get_tuple` now has to return `pending_writes` in `writes_sort_key` order, and memory, SQLite and Postgres do, so the default walk is right for any saver that follows that. `AsyncPostgresSaver` also gets the sync `get_delta_channel_history` that `AsyncSqliteSaver` already had; sync reads used to fall back to the default walk. Memory and Postgres already stored `task_path`. SQLite accepted it on `put_writes` and dropped it, so `writes` gains the column and `setup()` adds it to existing databases. Writes without a task path (graph input, `update_state` updates, rows written before the column existed) sort first by `task_id`, which is also where live execution applies input. ## Compatibility - Memory and Postgres already had `task_path` on disk, so existing threads replay in the corrected order after upgrade. - On SQLite, rows from before the upgrade read back as `''` and keep their old order. Only new writes get the corrected order, since there is no stored path to recover. - SQLite has no `ADD COLUMN IF NOT EXISTS`, so `setup()` runs the `ALTER` and treats `duplicate column name` as success, which also covers two processes migrating one file. That error comes while SQLite prepares the statement, before it asks for the write lock, so `setup()` on an up-to-date database doesn't wait on other writers. A read-only database from before the column is read with empty paths instead of failing setup. - Downgrade is safe: older code's `INSERT` omits the column and the column has a default. - `get_tuple` and `list` return `pending_writes` in `writes_sort_key` order instead of task id order (SQLite, Postgres) or put order (memory). Nothing in langgraph depends on the old order. - Postgres reads `task_path` and `idx` with each pending write, so a subclass that pairs `SELECT_SQL` with its own `_load_writes`, or the other way round, has to match the new row shape. - A third-party saver has to replay in `writes_sort_key` order: in its own `get_delta_channel_history` if it has one, otherwise by returning `pending_writes` from `get_tuple` in that order. Otherwise it fails the new conformance tests; the `get_tuple` one is a base test, so it runs for every saver. - langgraph, checkpoint-sqlite and checkpoint-postgres import `writes_sort_key` from langgraph-checkpoint, so this bumps langgraph-checkpoint to 4.3.0 and requires it there. 4.3.0 has to be released first. ## Limits - Several updates in one `bulk_update_state` super-step share a task id, so the saver keeps only the first one's writes and the `DeltaChannel` loses the rest (plain channels are fine). That's on main too; #9128 fixes it and stores a task path per update, so with this PR those updates also replay in the order given. - Exit durability stores its writes differently; resume ordering there is fixed in #9114, stacked on this PR. - langgraphjs applies `DeltaChannel` writes in task id order, live and on replay (langchain-ai/langgraphjs#2544), so its reads match too. Moving JS to task path order is a follow-up. ## Tests Each fails on main and passes here: - `test_delta_channel_parallel_order.py`, across memory, SQLite and Postgres: parallel writers and a `Send` fan-out, comparing `get_state` with the live result, including sync `get_state` on the async savers. - Three conformance tests whose task ids sort opposite to their task paths, so they cannot pass by chance: two on `get_delta_channel_history`, one on `get_tuple`'s `pending_writes` order. - SQLite migration tests from a pre-column database, sync and async, including a read-only file. One more that passes on main too: SQLite `setup()` while another connection holds the write lock, with no busy timeout, sync and async. It keeps the `ALTER` from waiting on that lock. Thanks to @ErenAta16 for the report, the reproduction, and for tracing `task_path` through all three backends, and to @ragnarok268 for the conformance-test design with task ids ordered opposite to their paths. --------- Co-authored-by: ErenAta16 <149434812+ErenAta16@users.noreply.github.com> Co-authored-by: ragnarok268 <58264829+ragnarok268@users.noreply.github.com>
756 lines
29 KiB
Python
756 lines
29 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import random
|
|
import threading
|
|
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any, TypeVar, cast
|
|
|
|
import aiosqlite
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.checkpoint.base import (
|
|
WRITES_IDX_MAP,
|
|
BaseCheckpointSaver,
|
|
ChannelVersions,
|
|
Checkpoint,
|
|
CheckpointMetadata,
|
|
CheckpointTuple,
|
|
DeltaChannelHistory,
|
|
SerializerProtocol,
|
|
get_checkpoint_id,
|
|
get_checkpoint_metadata,
|
|
)
|
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
|
|
|
from langgraph.checkpoint.sqlite._delta import (
|
|
DELTA_STAGE1_SQL,
|
|
build_delta_channels_writes_history,
|
|
build_delta_stage2_sql,
|
|
step_walk_with_row,
|
|
)
|
|
from langgraph.checkpoint.sqlite.utils import (
|
|
load_pending_writes,
|
|
pending_writes_sql,
|
|
search_where,
|
|
)
|
|
|
|
T = TypeVar("T", bound=Callable)
|
|
|
|
|
|
class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
"""An asynchronous checkpoint saver that stores checkpoints in a SQLite database.
|
|
|
|
This class provides an asynchronous interface for saving and retrieving checkpoints
|
|
using a SQLite database. It's designed for use in asynchronous environments and
|
|
offers better performance for I/O-bound operations compared to synchronous alternatives.
|
|
|
|
Attributes:
|
|
conn (aiosqlite.Connection): The asynchronous SQLite database connection.
|
|
serde (SerializerProtocol): The serializer used for encoding/decoding checkpoints.
|
|
|
|
Tip:
|
|
Requires the [aiosqlite](https://pypi.org/project/aiosqlite/) package.
|
|
Install it with `pip install aiosqlite`.
|
|
|
|
Warning:
|
|
While this class supports asynchronous checkpointing, it is not recommended
|
|
for production workloads due to limitations in SQLite's write performance.
|
|
For production use, consider a more robust database like PostgreSQL.
|
|
|
|
Tip:
|
|
Remember to **close the database connection** after executing your code,
|
|
otherwise, you may see the graph "hang" after execution (since the program
|
|
will not exit until the connection is closed).
|
|
|
|
The easiest way is to use the `async with` statement as shown in the examples.
|
|
|
|
```python
|
|
async with AsyncSqliteSaver.from_conn_string("checkpoints.sqlite") as saver:
|
|
# Your code here
|
|
graph = builder.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "thread-1"}}
|
|
async for event in graph.astream_events(..., config, version="v1"):
|
|
print(event)
|
|
```
|
|
|
|
Examples:
|
|
Usage within StateGraph:
|
|
|
|
```pycon
|
|
>>> import asyncio
|
|
>>>
|
|
>>> from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
>>> from langgraph.graph import StateGraph
|
|
>>>
|
|
>>> async def main():
|
|
>>> builder = StateGraph(int)
|
|
>>> builder.add_node("add_one", lambda x: x + 1)
|
|
>>> builder.set_entry_point("add_one")
|
|
>>> builder.set_finish_point("add_one")
|
|
>>> async with AsyncSqliteSaver.from_conn_string("checkpoints.db") as memory:
|
|
>>> graph = builder.compile(checkpointer=memory)
|
|
>>> coro = graph.ainvoke(1, {"configurable": {"thread_id": "thread-1"}})
|
|
>>> print(await asyncio.gather(coro))
|
|
>>>
|
|
>>> asyncio.run(main())
|
|
Output: [2]
|
|
```
|
|
Raw usage:
|
|
|
|
```pycon
|
|
>>> import asyncio
|
|
>>> import aiosqlite
|
|
>>> from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
>>>
|
|
>>> async def main():
|
|
>>> async with aiosqlite.connect("checkpoints.db") as conn:
|
|
... saver = AsyncSqliteSaver(conn)
|
|
... config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
|
|
... checkpoint = {"ts": "2023-05-03T10:00:00Z", "data": {"key": "value"}, "id": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}
|
|
... saved_config = await saver.aput(config, checkpoint, {}, {})
|
|
... print(saved_config)
|
|
>>> asyncio.run(main())
|
|
{'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '0c62ca34-ac19-445d-bbb0-5b4984975b2a'}}
|
|
```
|
|
"""
|
|
|
|
lock: asyncio.Lock
|
|
is_setup: bool
|
|
_has_task_path: bool = True
|
|
|
|
def __init__(
|
|
self,
|
|
conn: aiosqlite.Connection,
|
|
*,
|
|
serde: SerializerProtocol | None = None,
|
|
):
|
|
super().__init__(serde=serde)
|
|
self.jsonplus_serde = JsonPlusSerializer()
|
|
self.conn = conn
|
|
self.lock = asyncio.Lock()
|
|
self.loop = asyncio.get_running_loop()
|
|
self.is_setup = False
|
|
|
|
@classmethod
|
|
@asynccontextmanager
|
|
async def from_conn_string(
|
|
cls, conn_string: str
|
|
) -> AsyncIterator[AsyncSqliteSaver]:
|
|
"""Create a new AsyncSqliteSaver instance from a connection string.
|
|
|
|
Args:
|
|
conn_string: The SQLite connection string.
|
|
|
|
Yields:
|
|
AsyncSqliteSaver: A new AsyncSqliteSaver instance.
|
|
"""
|
|
async with aiosqlite.connect(conn_string) as conn:
|
|
yield cls(conn)
|
|
|
|
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
|
"""Get a checkpoint tuple from the database.
|
|
|
|
This method retrieves a checkpoint tuple from the SQLite database based on the
|
|
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
|
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
|
for the given thread ID is retrieved.
|
|
|
|
Args:
|
|
config: The config to use for retrieving the checkpoint.
|
|
|
|
Returns:
|
|
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
|
"""
|
|
try:
|
|
# check if we are in the main thread, only bg threads can block
|
|
# we don't check in other methods to avoid the overhead
|
|
if asyncio.get_running_loop() is self.loop:
|
|
raise asyncio.InvalidStateError(
|
|
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
|
"different thread. From the main thread, use the async interface. "
|
|
"For example, use `await checkpointer.aget_tuple(...)` or `await "
|
|
"graph.ainvoke(...)`."
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.aget_tuple(config), self.loop
|
|
).result()
|
|
|
|
def list(
|
|
self,
|
|
config: RunnableConfig | None,
|
|
*,
|
|
filter: dict[str, Any] | None = None,
|
|
before: RunnableConfig | None = None,
|
|
limit: int | None = None,
|
|
) -> Iterator[CheckpointTuple]:
|
|
"""List checkpoints from the database asynchronously.
|
|
|
|
This method retrieves a list of checkpoint tuples from the SQLite database based
|
|
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
|
|
|
Args:
|
|
config: Base configuration for filtering checkpoints.
|
|
filter: Additional filtering criteria for metadata.
|
|
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
limit: Maximum number of checkpoints to return.
|
|
|
|
Yields:
|
|
An iterator of matching checkpoint tuples.
|
|
"""
|
|
try:
|
|
# check if we are in the main thread, only bg threads can block
|
|
# we don't check in other methods to avoid the overhead
|
|
if asyncio.get_running_loop() is self.loop:
|
|
raise asyncio.InvalidStateError(
|
|
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
|
"different thread. From the main thread, use the async interface. "
|
|
"For example, use `checkpointer.alist(...)` or `await "
|
|
"graph.ainvoke(...)`."
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
|
|
while True:
|
|
try:
|
|
yield asyncio.run_coroutine_threadsafe(
|
|
anext(aiter_), # type: ignore[arg-type]
|
|
self.loop,
|
|
).result()
|
|
except StopAsyncIteration:
|
|
break
|
|
|
|
def put(
|
|
self,
|
|
config: RunnableConfig,
|
|
checkpoint: Checkpoint,
|
|
metadata: CheckpointMetadata,
|
|
new_versions: ChannelVersions,
|
|
) -> RunnableConfig:
|
|
"""Save a checkpoint to the database.
|
|
|
|
This method saves a checkpoint to the SQLite database. The checkpoint is associated
|
|
with the provided config and its parent config (if any).
|
|
|
|
Args:
|
|
config: The config to associate with the checkpoint.
|
|
checkpoint: The checkpoint to save.
|
|
metadata: Additional metadata to save with the checkpoint.
|
|
new_versions: New channel versions as of this write.
|
|
|
|
Returns:
|
|
RunnableConfig: Updated configuration after storing the checkpoint.
|
|
"""
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.aput(config, checkpoint, metadata, new_versions), self.loop
|
|
).result()
|
|
|
|
def put_writes(
|
|
self,
|
|
config: RunnableConfig,
|
|
writes: Sequence[tuple[str, Any]],
|
|
task_id: str,
|
|
task_path: str = "",
|
|
) -> None:
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.aput_writes(config, writes, task_id, task_path), self.loop
|
|
).result()
|
|
|
|
def delete_thread(self, thread_id: str) -> None:
|
|
"""Delete all checkpoints and writes associated with a thread ID.
|
|
|
|
Args:
|
|
thread_id: The thread ID to delete.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
try:
|
|
# check if we are in the main thread, only bg threads can block
|
|
# we don't check in other methods to avoid the overhead
|
|
if asyncio.get_running_loop() is self.loop:
|
|
raise asyncio.InvalidStateError(
|
|
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
|
"different thread. From the main thread, use the async interface. "
|
|
"For example, use `checkpointer.alist(...)` or `await "
|
|
"graph.ainvoke(...)`."
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.adelete_thread(thread_id), self.loop
|
|
).result()
|
|
|
|
def get_delta_channel_history(
|
|
self, *, config: RunnableConfig, channels: Sequence[str]
|
|
) -> Mapping[str, DeltaChannelHistory]:
|
|
"""Sync bridge to `aget_delta_channel_history`.
|
|
|
|
Mirrors the same cross-thread guard as `get_tuple` /
|
|
`delete_thread` — calling from the loop thread raises rather than
|
|
deadlocking.
|
|
"""
|
|
try:
|
|
if asyncio.get_running_loop() is self.loop:
|
|
raise asyncio.InvalidStateError(
|
|
"Synchronous calls to AsyncSqliteSaver are only allowed from a "
|
|
"different thread. From the main thread, use the async interface. "
|
|
"For example, use `await checkpointer.aget_delta_channel_history(...)`."
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.aget_delta_channel_history(config=config, channels=channels),
|
|
self.loop,
|
|
).result()
|
|
|
|
async def setup(self) -> None:
|
|
"""Set up the checkpoint database asynchronously.
|
|
|
|
This method creates the necessary tables in the SQLite database if they don't
|
|
already exist. It is called automatically when needed and should not be called
|
|
directly by the user.
|
|
"""
|
|
async with self.lock:
|
|
if self.is_setup:
|
|
return
|
|
await _ensure_connected(self.conn)
|
|
async with self.conn.executescript(
|
|
"""
|
|
PRAGMA journal_mode=WAL;
|
|
CREATE TABLE IF NOT EXISTS checkpoints (
|
|
thread_id TEXT NOT NULL,
|
|
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
|
checkpoint_id TEXT NOT NULL,
|
|
parent_checkpoint_id TEXT,
|
|
type TEXT,
|
|
checkpoint BLOB,
|
|
metadata BLOB,
|
|
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS writes (
|
|
thread_id TEXT NOT NULL,
|
|
checkpoint_ns TEXT NOT NULL DEFAULT '',
|
|
checkpoint_id TEXT NOT NULL,
|
|
task_id TEXT NOT NULL,
|
|
task_path TEXT NOT NULL DEFAULT '',
|
|
idx INTEGER NOT NULL,
|
|
channel TEXT NOT NULL,
|
|
type TEXT,
|
|
value BLOB,
|
|
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
|
);
|
|
"""
|
|
):
|
|
await self.conn.commit()
|
|
|
|
# sqlite has no ADD COLUMN IF NOT EXISTS; this migrates databases
|
|
# created before `task_path` existed and is a no-op on the rest.
|
|
try:
|
|
await self.conn.execute(
|
|
"ALTER TABLE writes ADD COLUMN task_path TEXT NOT NULL DEFAULT ''"
|
|
)
|
|
await self.conn.commit()
|
|
except aiosqlite.OperationalError as e:
|
|
# A read-only database from before the column can still be read;
|
|
# its rows would all read back as '' anyway.
|
|
if "readonly database" in str(e):
|
|
self._has_task_path = False
|
|
elif "duplicate column name" not in str(e):
|
|
raise
|
|
|
|
self.is_setup = True
|
|
|
|
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
|
"""Get a checkpoint tuple from the database asynchronously.
|
|
|
|
This method retrieves a checkpoint tuple from the SQLite database based on the
|
|
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
|
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
|
for the given thread ID is retrieved.
|
|
|
|
Args:
|
|
config: The config to use for retrieving the checkpoint.
|
|
|
|
Returns:
|
|
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
|
"""
|
|
await self.setup()
|
|
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
|
async with self.lock, self.conn.cursor() as cur:
|
|
# find the latest checkpoint for the thread_id
|
|
if checkpoint_id := get_checkpoint_id(config):
|
|
await cur.execute(
|
|
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
|
|
(
|
|
str(config["configurable"]["thread_id"]),
|
|
checkpoint_ns,
|
|
checkpoint_id,
|
|
),
|
|
)
|
|
else:
|
|
await cur.execute(
|
|
"SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1",
|
|
(str(config["configurable"]["thread_id"]), checkpoint_ns),
|
|
)
|
|
# if a checkpoint is found, return it
|
|
if value := await cur.fetchone():
|
|
(
|
|
thread_id,
|
|
checkpoint_id,
|
|
parent_checkpoint_id,
|
|
type,
|
|
checkpoint,
|
|
metadata,
|
|
) = value
|
|
if not get_checkpoint_id(config):
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": checkpoint_id,
|
|
}
|
|
}
|
|
# find any pending writes
|
|
await cur.execute(
|
|
pending_writes_sql(self._has_task_path),
|
|
(
|
|
str(config["configurable"]["thread_id"]),
|
|
checkpoint_ns,
|
|
str(config["configurable"]["checkpoint_id"]),
|
|
),
|
|
)
|
|
# deserialize the checkpoint and metadata
|
|
return CheckpointTuple(
|
|
config,
|
|
self.serde.loads_typed((type, checkpoint)),
|
|
cast(
|
|
CheckpointMetadata,
|
|
(json.loads(metadata) if metadata is not None else {}),
|
|
),
|
|
(
|
|
{
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": parent_checkpoint_id,
|
|
}
|
|
}
|
|
if parent_checkpoint_id
|
|
else None
|
|
),
|
|
load_pending_writes(await cur.fetchall(), self.serde),
|
|
)
|
|
|
|
async def alist(
|
|
self,
|
|
config: RunnableConfig | None,
|
|
*,
|
|
filter: dict[str, Any] | None = None,
|
|
before: RunnableConfig | None = None,
|
|
limit: int | None = None,
|
|
) -> AsyncIterator[CheckpointTuple]:
|
|
"""List checkpoints from the database asynchronously.
|
|
|
|
This method retrieves a list of checkpoint tuples from the SQLite database based
|
|
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
|
|
|
|
Args:
|
|
config: Base configuration for filtering checkpoints.
|
|
filter: Additional filtering criteria for metadata.
|
|
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
limit: Maximum number of checkpoints to return.
|
|
|
|
Yields:
|
|
An asynchronous iterator of matching checkpoint tuples.
|
|
"""
|
|
await self.setup()
|
|
where, params = search_where(config, filter, before)
|
|
query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
|
|
FROM checkpoints
|
|
{where}
|
|
ORDER BY checkpoint_id DESC"""
|
|
if limit is not None:
|
|
query += " LIMIT ?"
|
|
params = (*params, limit)
|
|
async with (
|
|
self.lock,
|
|
self.conn.execute(query, params) as cur,
|
|
self.conn.cursor() as wcur,
|
|
):
|
|
async for (
|
|
thread_id,
|
|
checkpoint_ns,
|
|
checkpoint_id,
|
|
parent_checkpoint_id,
|
|
type,
|
|
checkpoint,
|
|
metadata,
|
|
) in cur:
|
|
await wcur.execute(
|
|
pending_writes_sql(self._has_task_path),
|
|
(thread_id, checkpoint_ns, checkpoint_id),
|
|
)
|
|
yield CheckpointTuple(
|
|
{
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": checkpoint_id,
|
|
}
|
|
},
|
|
self.serde.loads_typed((type, checkpoint)),
|
|
cast(
|
|
CheckpointMetadata,
|
|
(json.loads(metadata) if metadata is not None else {}),
|
|
),
|
|
(
|
|
{
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": parent_checkpoint_id,
|
|
}
|
|
}
|
|
if parent_checkpoint_id
|
|
else None
|
|
),
|
|
load_pending_writes(await wcur.fetchall(), self.serde),
|
|
)
|
|
|
|
async def aput(
|
|
self,
|
|
config: RunnableConfig,
|
|
checkpoint: Checkpoint,
|
|
metadata: CheckpointMetadata,
|
|
new_versions: ChannelVersions,
|
|
) -> RunnableConfig:
|
|
"""Save a checkpoint to the database asynchronously.
|
|
|
|
This method saves a checkpoint to the SQLite database. The checkpoint is associated
|
|
with the provided config and its parent config (if any).
|
|
|
|
Args:
|
|
config: The config to associate with the checkpoint.
|
|
checkpoint: The checkpoint to save.
|
|
metadata: Additional metadata to save with the checkpoint.
|
|
new_versions: New channel versions as of this write.
|
|
|
|
Returns:
|
|
RunnableConfig: Updated configuration after storing the checkpoint.
|
|
"""
|
|
await self.setup()
|
|
thread_id = config["configurable"]["thread_id"]
|
|
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
|
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
|
|
serialized_metadata = json.dumps(
|
|
get_checkpoint_metadata(config, metadata), ensure_ascii=False
|
|
).encode("utf-8", "ignore")
|
|
async with (
|
|
self.lock,
|
|
self.conn.execute(
|
|
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
|
(
|
|
str(config["configurable"]["thread_id"]),
|
|
checkpoint_ns,
|
|
checkpoint["id"],
|
|
config["configurable"].get("checkpoint_id"),
|
|
type_,
|
|
serialized_checkpoint,
|
|
serialized_metadata,
|
|
),
|
|
),
|
|
):
|
|
await self.conn.commit()
|
|
return {
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": checkpoint["id"],
|
|
}
|
|
}
|
|
|
|
async def aput_writes(
|
|
self,
|
|
config: RunnableConfig,
|
|
writes: Sequence[tuple[str, Any]],
|
|
task_id: str,
|
|
task_path: str = "",
|
|
) -> None:
|
|
"""Store intermediate writes linked to a checkpoint asynchronously.
|
|
|
|
This method saves intermediate writes associated with a checkpoint to the database.
|
|
|
|
Args:
|
|
config: Configuration of the related checkpoint.
|
|
writes: List of writes to store, each as (channel, value) pair.
|
|
task_id: Identifier for the task creating the writes.
|
|
task_path: Path of the task creating the writes.
|
|
"""
|
|
query = (
|
|
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
|
if all(w[0] in WRITES_IDX_MAP for w in writes)
|
|
else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
|
)
|
|
await self.setup()
|
|
async with self.lock, self.conn.cursor() as cur:
|
|
await cur.executemany(
|
|
query,
|
|
[
|
|
(
|
|
str(config["configurable"]["thread_id"]),
|
|
str(config["configurable"]["checkpoint_ns"]),
|
|
str(config["configurable"]["checkpoint_id"]),
|
|
task_id,
|
|
task_path,
|
|
WRITES_IDX_MAP.get(channel, idx),
|
|
channel,
|
|
*self.serde.dumps_typed(value),
|
|
)
|
|
for idx, (channel, value) in enumerate(writes)
|
|
],
|
|
)
|
|
await self.conn.commit()
|
|
|
|
async def adelete_thread(self, thread_id: str) -> None:
|
|
"""Delete all checkpoints and writes associated with a thread ID.
|
|
|
|
Args:
|
|
thread_id: The thread ID to delete.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
async with self.lock, self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"DELETE FROM checkpoints WHERE thread_id = ?",
|
|
(str(thread_id),),
|
|
)
|
|
await cur.execute(
|
|
"DELETE FROM writes WHERE thread_id = ?",
|
|
(str(thread_id),),
|
|
)
|
|
await self.conn.commit()
|
|
|
|
async def aget_delta_channel_history(
|
|
self, *, config: RunnableConfig, channels: Sequence[str]
|
|
) -> Mapping[str, DeltaChannelHistory]:
|
|
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
|
|
|
|
See `SqliteSaver.get_delta_channel_history` for design notes; this
|
|
is the async equivalent using `aiosqlite` cursors. Stage 1 streams
|
|
the parent chain from the target and Python-deserializes each
|
|
checkpoint blob to find per-channel snapshots; stage 2 fetches
|
|
only the relevant writes via per-channel UNION ALL.
|
|
"""
|
|
if not channels:
|
|
return {}
|
|
channels = list(channels)
|
|
await self.setup()
|
|
thread_id = str(config["configurable"]["thread_id"])
|
|
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
|
checkpoint_id = get_checkpoint_id(config)
|
|
if checkpoint_id is None:
|
|
target = await self.aget_tuple(config)
|
|
if target is None:
|
|
return {ch: {"writes": []} for ch in channels}
|
|
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
|
|
|
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
|
seed_val_by_ch: dict[str, Any] = {}
|
|
walk_state: dict[str, Any] = {}
|
|
seeded: set[str] = set()
|
|
|
|
async with self.lock, self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
DELTA_STAGE1_SQL,
|
|
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
|
|
)
|
|
async for row in cur:
|
|
cid, type_tag, blob = row
|
|
if step_walk_with_row(
|
|
cid=cid,
|
|
type_tag=type_tag,
|
|
blob=blob,
|
|
target_id=checkpoint_id,
|
|
serde=self.serde,
|
|
chain_by_ch=chain_by_ch,
|
|
seed_val_by_ch=seed_val_by_ch,
|
|
walk_state=walk_state,
|
|
seeded=seeded,
|
|
channels=channels,
|
|
):
|
|
break
|
|
|
|
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
|
stage2_sql = build_delta_stage2_sql(
|
|
has_task_path=self._has_task_path,
|
|
chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain],
|
|
)
|
|
if stage2_sql:
|
|
stage2_params: list[Any] = []
|
|
for ch in channels_with_chain:
|
|
stage2_params.extend(
|
|
[thread_id, checkpoint_ns, ch, *chain_by_ch[ch]]
|
|
)
|
|
await cur.execute(stage2_sql, stage2_params)
|
|
stage2_rows = cast(
|
|
"list[tuple[str, str, str, int, str, bytes, str]]",
|
|
await cur.fetchall(),
|
|
)
|
|
else:
|
|
stage2_rows = []
|
|
|
|
return build_delta_channels_writes_history(
|
|
channels=channels,
|
|
chain_by_ch=chain_by_ch,
|
|
seed_val_by_ch=seed_val_by_ch,
|
|
seeded=seeded,
|
|
stage2_rows=stage2_rows,
|
|
serde=self.serde,
|
|
)
|
|
|
|
def get_next_version(self, current: str | None, channel: None) -> str:
|
|
"""Generate the next version ID for a channel.
|
|
|
|
This method creates a new version identifier for a channel based on its current version.
|
|
|
|
Args:
|
|
current (Optional[str]): The current version identifier of the channel.
|
|
|
|
Returns:
|
|
str: The next version identifier, which is guaranteed to be monotonically increasing.
|
|
"""
|
|
if current is None:
|
|
current_v = 0
|
|
elif isinstance(current, int):
|
|
current_v = current
|
|
else:
|
|
current_v = int(current.split(".")[0])
|
|
next_v = current_v + 1
|
|
next_h = random.random()
|
|
return f"{next_v:032}.{next_h:016}"
|
|
|
|
|
|
async def _ensure_connected(conn: aiosqlite.Connection) -> None:
|
|
if not _CONN_STARTED_CHECK(conn):
|
|
await conn
|
|
|
|
|
|
def _build_conn_started_check() -> Callable[[aiosqlite.Connection], bool]:
|
|
is_alive = getattr(aiosqlite.Connection, "is_alive", None)
|
|
if callable(is_alive):
|
|
return lambda conn: conn.is_alive() # type: ignore[attr-defined]
|
|
|
|
def _started(conn: aiosqlite.Connection) -> bool:
|
|
thread: threading.Thread | None = getattr(conn, "_thread", None)
|
|
return False if thread is None else thread.is_alive()
|
|
|
|
return _started
|
|
|
|
|
|
_CONN_STARTED_CHECK = _build_conn_started_check()
|