From 18a9ae45f311743e21c74732fa0daa614b0d4c65 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 17 Apr 2025 09:38:58 -0700 Subject: [PATCH] Add delete_thread method to Checkpointer class (#4328) - Deletes all data associated with a thread_id - Implemented in InMemory, Sqlite and Postgres checkpointers Co-authored-by: Eugene Yurtsev --- .../langgraph/checkpoint/postgres/__init__.py | 23 +++++++++ .../langgraph/checkpoint/postgres/aio.py | 48 +++++++++++++++++++ .../langgraph/checkpoint/sqlite/__init__.py | 19 ++++++++ .../langgraph/checkpoint/sqlite/aio.py | 45 +++++++++++++++++ .../langgraph/checkpoint/base/__init__.py | 22 +++++++++ .../langgraph/checkpoint/memory/__init__.py | 33 ++++++++++++- libs/langgraph/tests/test_pregel.py | 15 +++++- 7 files changed, 201 insertions(+), 4 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 1c3426603..f4d8a322a 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -357,6 +357,29 @@ class PostgresSaver(BasePostgresSaver): ), ) + def delete_thread(self, thread_id: str) -> None: + """Delete all checkpoints and writes associated with a thread ID. + + Args: + thread_id (str): The thread ID to delete. + + Returns: + None + """ + with self._cursor(pipeline=True) as cur: + cur.execute( + "DELETE FROM checkpoints WHERE thread_id = %s", + (str(thread_id),), + ) + cur.execute( + "DELETE FROM checkpoint_blobs WHERE thread_id = %s", + (str(thread_id),), + ) + cur.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = %s", + (str(thread_id),), + ) + @contextmanager def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: """Create a database cursor as a context manager. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index a19b7bfc1..06354bb8f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -314,6 +314,29 @@ class AsyncPostgresSaver(BasePostgresSaver): async with self._cursor(pipeline=True) as cur: await cur.executemany(query, params) + async def adelete_thread(self, thread_id: str) -> None: + """Delete all checkpoints and writes associated with a thread ID. + + Args: + thread_id (str): The thread ID to delete. + + Returns: + None + """ + async with self._cursor(pipeline=True) as cur: + await cur.execute( + "DELETE FROM checkpoints WHERE thread_id = %s", + (str(thread_id),), + ) + await cur.execute( + "DELETE FROM checkpoint_blobs WHERE thread_id = %s", + (str(thread_id),), + ) + await cur.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = %s", + (str(thread_id),), + ) + @asynccontextmanager async def _cursor( self, *, pipeline: bool = False @@ -481,5 +504,30 @@ class AsyncPostgresSaver(BasePostgresSaver): 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 (str): 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 AsyncPostgresSaver 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.adelete_thread(thread_id), self.loop + ).result() + __all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"] diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index a53728ab1..427a38341 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -464,6 +464,25 @@ class SqliteSaver(BaseCheckpointSaver[str]): ], ) + def delete_thread(self, thread_id: str) -> None: + """Delete all checkpoints and writes associated with a thread ID. + + Args: + thread_id (str): The thread ID to delete. + + Returns: + None + """ + with self.cursor() as cur: + cur.execute( + "DELETE FROM checkpoints WHERE thread_id = ?", + (str(thread_id),), + ) + cur.execute( + "DELETE FROM writes WHERE thread_id = ?", + (str(thread_id),), + ) + async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database asynchronously. diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 4774624e1..574bfafda 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -244,6 +244,31 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): 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 (str): 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() + async def setup(self) -> None: """Set up the checkpoint database asynchronously. @@ -535,6 +560,26 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): ) 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 (str): 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() + def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: """Generate the next version ID for a channel. diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index fbaa45fe4..35bc79bdc 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -321,6 +321,17 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + def delete_thread( + self, + thread_id: str, + ) -> None: + """Delete all checkpoints and writes associated with a specific thread ID. + + Args: + thread_id (str): The thread ID whose checkpoints should be deleted. + """ + raise NotImplementedError + async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: """Asynchronously fetch a checkpoint using the given configuration. @@ -415,6 +426,17 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + async def adelete_thread( + self, + thread_id: str, + ) -> None: + """Delete all checkpoints and writes associated with a specific thread ID. + + Args: + thread_id (str): The thread ID whose checkpoints should be deleted. + """ + raise NotImplementedError + def get_next_version(self, current: Optional[V], channel: ChannelProtocol) -> V: """Generate the next version ID for a channel. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 863048240..f73b5c746 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -69,7 +69,7 @@ class InMemorySaver( ], ] writes: defaultdict[ - tuple[str, str, str], + tuple[str, str, str], # thread ID, checkpoint NS, checkpoint ID dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]], ] blobs: dict[ @@ -451,6 +451,24 @@ class InMemorySaver( task_path, ) + def delete_thread(self, thread_id: str) -> None: + """Delete all checkpoints and writes associated with a thread ID. + + Args: + thread_id (str): The thread ID to delete. + + Returns: + None + """ + if thread_id in self.storage: + del self.storage[thread_id] + for k in list(self.writes.keys()): + if k[0] == thread_id: + del self.writes[k] + for k in list(self.blobs.keys()): + if k[0] == thread_id: + del self.blobs[k] + async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Asynchronous version of get_tuple. @@ -530,6 +548,17 @@ class InMemorySaver( """ return self.put_writes(config, writes, task_id, task_path) + async def adelete_thread(self, thread_id: str) -> None: + """Delete all checkpoints and writes associated with a thread ID. + + Args: + thread_id (str): The thread ID to delete. + + Returns: + None + """ + return self.delete_thread(thread_id) + def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str: if current is None: current_v = 0 @@ -615,4 +644,4 @@ class PersistentDict(defaultdict): except Exception: logging.error(f"Failed to load file: {fileobj.name}") raise - raise ValueError("File not in a supported f ormat") + raise ValueError("File not in a supported format") diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 27f05c074..c458c7dbc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4601,7 +4601,7 @@ def test_checkpoint_metadata() -> None: assert chkpnt_tuple.metadata["test_config_4"] == "bar" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_remove_message_via_state_update( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: @@ -4631,6 +4631,12 @@ def test_remove_message_via_state_update( assert len(updated_state.values) == 1 assert updated_state.values[-1].content == "Hi" + app.checkpointer.delete_thread(config["configurable"]["thread_id"]) + + # Verify that the message was removed from the checkpointer + assert app.checkpointer.get_tuple(config) is None + assert [*app.get_state_history(config)] == [] + def test_remove_message_from_node(): from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage @@ -6034,7 +6040,7 @@ def test_concurrent_execution_thread_safety(): assert result["counter"] == 1 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: str): """Test recovery from checkpoints after failures.""" checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -6087,6 +6093,11 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error) assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error + # Verify delete leaves it empty + graph.checkpointer.delete_thread(config["configurable"]["thread_id"]) + assert graph.checkpointer.get_tuple(config) is None + assert [*graph.get_state_history(config)] == [] + def test_multiple_updates_root() -> None: def node_a(state):