Update sqlite signature

This commit is contained in:
Nuno Campos
2025-01-15 10:58:46 -08:00
parent bba00506ea
commit eed577ee2a
2 changed files with 22 additions and 13 deletions
@@ -424,6 +424,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint.
@@ -433,6 +434,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
"""
query = (
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
@@ -398,9 +398,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
async with self.lock, self.conn.execute(
query, params
) as cur, self.conn.cursor() as wcur:
async with (
self.lock,
self.conn.execute(query, params) as cur,
self.conn.cursor() as wcur,
):
async for (
thread_id,
checkpoint_ns,
@@ -467,16 +469,19 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
checkpoint_ns = config["configurable"]["checkpoint_ns"]
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
serialized_metadata = self.jsonplus_serde.dumps(metadata)
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,
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()
@@ -493,6 +498,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
config: RunnableConfig,
writes: Sequence[Tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
@@ -502,6 +508,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
task_path (str): Path of the task creating the writes.
"""
query = (
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"