From 7aaeedd7ba114927016c0dbd3798a51197520466 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 19 Aug 2024 15:40:04 -0700 Subject: [PATCH 1/2] checkpoint-*: In put_writes clear any previously saved writes for this task_id - Previously implementations were clearing only tasks if the index matched a previously saved one, which isn't enough to guarantee idempotency --- .../langgraph/checkpoint/postgres/__init__.py | 8 +++++ .../langgraph/checkpoint/postgres/aio.py | 8 +++++ .../langgraph/checkpoint/postgres/base.py | 11 +++++++ .../langgraph/checkpoint/sqlite/__init__.py | 10 ++++++ .../langgraph/checkpoint/sqlite/aio.py | 32 ++++++++++++------- .../langgraph/checkpoint/memory/__init__.py | 1 + 6 files changed, 59 insertions(+), 11 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index e07f0aa5f..153ba4bc3 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -317,6 +317,14 @@ class PostgresSaver(BasePostgresSaver): task_id (str): Identifier for the task creating the writes. """ with self._cursor() as cur: + cur.execute( + self.DELETE_WRITES_SQL, + config["configurable"]["thread_id"], + config["configurable"]["checkpoint_ns"], + config["configurable"]["checkpoint_id"], + task_id, + len(writes), + ) cur.executemany( self.UPSERT_CHECKPOINT_WRITES_SQL, self._dump_writes( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index c325399e2..6862df5bd 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -273,6 +273,14 @@ class AsyncPostgresSaver(BasePostgresSaver): task_id (str): Identifier for the task creating the writes. """ async with self._cursor() as cur: + await cur.execute( + self.DELETE_WRITES_SQL, + config["configurable"]["thread_id"], + config["configurable"]["checkpoint_ns"], + config["configurable"]["checkpoint_id"], + task_id, + len(writes), + ) await cur.executemany( self.UPSERT_CHECKPOINT_WRITES_SQL, await asyncio.to_thread( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index e9ca12305..506a13a27 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -105,6 +105,15 @@ UPSERT_CHECKPOINT_WRITES_SQL = """ ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING """ +DELETE_WRITES_SQL = """ + DELETE FROM checkpoint_writes + WHERE thread_id = %s + AND checkpoint_ns = %s + AND checkpoint_id = %s + AND task_id = %s + AND idx >= %s +""" + class BasePostgresSaver(BaseCheckpointSaver): SELECT_SQL = SELECT_SQL @@ -112,6 +121,8 @@ class BasePostgresSaver(BaseCheckpointSaver): UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL + DELETE_WRITES_SQL = DELETE_WRITES_SQL + jsonplus_serde = JsonPlusSerializer() def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint: diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index ac2f64f32..9bf4139ad 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -424,6 +424,16 @@ 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 (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 95b77209a..56d14f613 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -432,19 +432,29 @@ class AsyncSqliteSaver(BaseCheckpointSaver): task_id (str): Identifier for the task creating the writes. """ await self.setup() - async with self.conn.executemany( - "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", - [ + async with self.conn.cursor() as cur: + await 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, - idx, - channel, - *self.serde.dumps_typed(value), - ) - for idx, (channel, value) in enumerate(writes) - ], - ): - await self.conn.commit() + len(writes), + ), + ) + await cur.executemany( + "INSERT OR REPLACE 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, + channel, + *self.serde.dumps_typed(value), + ) + for idx, (channel, value) in enumerate(writes) + ], + ) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 5e59accd7..3e8d328e6 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -288,6 +288,7 @@ class MemorySaver( checkpoint_ns = config["configurable"]["checkpoint_ns"] checkpoint_id = config["configurable"]["checkpoint_id"] key = (thread_id, checkpoint_ns, checkpoint_id) + self.writes[key] = [w for w in self.writes[key] if w[0] != task_id] self.writes[key].extend( [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes] ) From c0c534cbe3c546329d2bd44814565a831da04c2f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 19 Aug 2024 15:44:22 -0700 Subject: [PATCH 2/2] Add pipeline arg --- .../langgraph/checkpoint/postgres/__init__.py | 2 +- libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 153ba4bc3..767b8f0e8 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -316,7 +316,7 @@ class PostgresSaver(BasePostgresSaver): writes (List[Tuple[str, Any]]): List of writes to store. task_id (str): Identifier for the task creating the writes. """ - with self._cursor() as cur: + with self._cursor(pipeline=True) as cur: cur.execute( self.DELETE_WRITES_SQL, config["configurable"]["thread_id"], diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 6862df5bd..cdbbe990a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -272,7 +272,7 @@ class AsyncPostgresSaver(BasePostgresSaver): 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. """ - async with self._cursor() as cur: + async with self._cursor(pipeline=True) as cur: await cur.execute( self.DELETE_WRITES_SQL, config["configurable"]["thread_id"],