diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index e07f0aa5f..767b8f0e8 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -316,7 +316,15 @@ 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"], + 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..cdbbe990a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -272,7 +272,15 @@ 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"], + 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] )