Fix semantics of put_writes/list (#1436)

* Fix semantics of put_writes/list

- put_writes(error) should not prevent saving future successful if task is retried successfully
- put_writes(writes) should be a no-op if non-error writes already exist for that task (this prevents tasks executed more than once from modifying writes previously saved / acted on)
- checkpoints should not include channel default values (ie those without a version)
- list() should fetch and return writes for each checkpoint

* Lint

* Rm print

* Fix import

* Lint
This commit is contained in:
Nuno Campos
2024-08-22 19:04:42 +00:00
committed by GitHub
parent 38daba5259
commit a261e1a497
14 changed files with 480 additions and 343 deletions
@@ -1,12 +1,13 @@
import sqlite3
import threading
from contextlib import contextmanager
from contextlib import closing, contextmanager
from hashlib import md5
from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
@@ -318,7 +319,7 @@ class SqliteSaver(BaseCheckpointSaver):
ORDER BY checkpoint_id DESC"""
if limit:
query += f" LIMIT {limit}"
with self.cursor(transaction=False) as cur:
with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
cur.execute(query, param_values)
for (
thread_id,
@@ -329,6 +330,10 @@ class SqliteSaver(BaseCheckpointSaver):
checkpoint,
metadata,
) in cur:
wcur.execute(
"SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
(thread_id, checkpoint_ns, checkpoint_id),
)
yield CheckpointTuple(
{
"configurable": {
@@ -350,6 +355,10 @@ class SqliteSaver(BaseCheckpointSaver):
if parent_checkpoint_id
else None
),
[
(task_id, channel, self.serde.loads_typed((type, value)))
for task_id, channel, type, value in wcur
],
)
def put(
@@ -424,25 +433,15 @@ class SqliteSaver(BaseCheckpointSaver):
task_id (str): Identifier for the task creating the writes.
"""
with self.lock, self.cursor() as cur:
cur.execute(
"DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
len(writes),
),
)
cur.executemany(
"INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
"INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
[
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["checkpoint_ns"]),
str(config["configurable"]["checkpoint_id"]),
task_id,
idx,
WRITES_IDX_MAP.get(channel, idx),
channel,
*self.serde.dumps_typed(value),
)