From 29d7a812ae238bf008064de3ecefd51cb36bf83f Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Fri, 19 Jul 2024 17:52:47 -0700 Subject: [PATCH] [Docs] Update checkpointer docstrings (#1074) --- docs/docs/reference/checkpoints.md | 4 -- .../langgraph/checkpoint/aiosqlite.py | 49 ++++++++++++------- libs/langgraph/langgraph/checkpoint/base.py | 21 ++++++-- libs/langgraph/langgraph/checkpoint/memory.py | 33 +++++++++++-- libs/langgraph/langgraph/checkpoint/sqlite.py | 21 ++++++++ 5 files changed, 99 insertions(+), 29 deletions(-) diff --git a/docs/docs/reference/checkpoints.md b/docs/docs/reference/checkpoints.md index 34c7dbb58..eb8735dd6 100644 --- a/docs/docs/reference/checkpoints.md +++ b/docs/docs/reference/checkpoints.md @@ -18,12 +18,10 @@ You can [compile][langgraph.graph.MessageGraph.compile] any LangGraph workflow w ### BaseCheckpointSaver ::: langgraph.checkpoint.base.BaseCheckpointSaver -handler: python ### SerializerProtocol ::: langgraph.checkpoint.SerializerProtocol -handler: python ## Implementations @@ -32,12 +30,10 @@ LangGraph also natively provides the following checkpoint implementations. ### MemorySaver ::: langgraph.checkpoint.memory.MemorySaver -handler: python ### AsyncSqliteSaver ::: langgraph.checkpoint.aiosqlite.AsyncSqliteSaver -handler: python ### SqliteSaver diff --git a/libs/langgraph/langgraph/checkpoint/aiosqlite.py b/libs/langgraph/langgraph/checkpoint/aiosqlite.py index 431684d1a..76ab9bde3 100644 --- a/libs/langgraph/langgraph/checkpoint/aiosqlite.py +++ b/libs/langgraph/langgraph/checkpoint/aiosqlite.py @@ -46,22 +46,29 @@ def not_implemented_sync_method(func: T) -> T: class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): """An asynchronous checkpoint saver that stores checkpoints in a SQLite database. + This class provides an asynchronous interface for saving and retrieving checkpoints + using a SQLite database. It's designed for use in asynchronous environments and + offers better performance for I/O-bound operations compared to synchronous alternatives. + + Attributes: + conn (aiosqlite.Connection): The asynchronous SQLite database connection. + serde (SerializerProtocol): The serializer used for encoding/decoding checkpoints. + Tip: Requires the [aiosqlite](https://pypi.org/project/aiosqlite/) package. Install it with `pip install aiosqlite`. - Note: - While this class does support asynchronous checkpointing, it is not recommended - for production workloads, due to limitations in SQLite's write performance. For - production workloads, consider using a more robust database like PostgreSQL. + Warning: + While this class supports asynchronous checkpointing, it is not recommended + for production workloads due to limitations in SQLite's write performance. + For production use, consider a more robust database like PostgreSQL. - !!! Important + Tip: Remember to **close the database connection** after executing your code, otherwise, you may see the graph "hang" after execution (since the program will not exit until the connection is closed). - The easiest way to do this is to use the `async with` statement, as shown in the - examples below. + The easiest way is to use the `async with` statement as shown in the examples. ```python async with AsyncSqliteSaver.from_conn_string("checkpoints.sqlite") as saver: @@ -72,12 +79,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): print(event) ``` - Args: - conn (aiosqlite.Connection): The asynchronous SQLite database connection. - serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to JsonPlusSerializerCompat. - Examples: - Usage within a StateGraph: + Usage within StateGraph: + ```pycon >>> import asyncio >>> import aiosqlite @@ -95,8 +99,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): >>> asyncio.run(coro) Output: 2 ``` - Raw usage: + ```pycon >>> import asyncio >>> import aiosqlite @@ -309,12 +313,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): on the provided config. The checkpoints are ordered by timestamp in descending order. Args: - config (RunnableConfig): The config to use for listing the checkpoints. - before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None. - limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None. + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): List checkpoints created before this configuration. + limit (Optional[int]): Maximum number of checkpoints to return. Yields: - AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples. + AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples. """ await self.setup() where, param_values = search_where(config, filter, before) @@ -356,6 +361,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): Args: config (RunnableConfig): The config to associate with the checkpoint. checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. Returns: RunnableConfig: The updated config containing the saved checkpoint's timestamp. @@ -385,6 +391,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): writes: Sequence[Tuple[str, Any]], task_id: str, ) -> None: + """Store intermediate writes linked to a checkpoint asynchronously. + + This method saves intermediate writes associated with a checkpoint to the database. + + Args: + 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. + """ await self.setup() async with self.conn.executemany( "INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)", diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 7f1ffab1f..b6694db89 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -29,6 +29,8 @@ PendingWrite = Tuple[str, str, Any] # Marked as total=False to allow for future expansion. class CheckpointMetadata(TypedDict, total=False): + """Metadata associated with a checkpoint.""" + source: Literal["input", "loop", "update"] """The source of the checkpoint. - "input": The checkpoint was created from an input to invoke/stream/batch. @@ -119,6 +121,8 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: class CheckpointTuple(NamedTuple): + """A tuple containing a checkpoint and its associated data.""" + config: RunnableConfig checkpoint: Checkpoint metadata: CheckpointMetadata @@ -269,11 +273,13 @@ class BaseCheckpointSaver(ABC): ) async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: - """ - Asynchronously fetch a checkpoint using the given configuration. + """Asynchronously fetch a checkpoint using the given configuration. Args: config (RunnableConfig): Configuration specifying which checkpoint to retrieve. + + Returns: + Optional[Checkpoint]: The requested checkpoint, or None if not found. """ if value := await self.aget_tuple(config): return value.checkpoint @@ -286,6 +292,9 @@ class BaseCheckpointSaver(ABC): Returns: Optional[CheckpointTuple]: The requested checkpoint tuple, or None if not found. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. """ raise NotImplementedError @@ -301,12 +310,15 @@ class BaseCheckpointSaver(ABC): Args: config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. - filter (Optional[Dict[str, Any]]): Additional filtering criteria. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. before (Optional[RunnableConfig]): List checkpoints created before this configuration. limit (Optional[int]): Maximum number of checkpoints to return. Returns: AsyncIterator[CheckpointTuple]: Async iterator of matching checkpoint tuples. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. """ raise NotImplementedError yield @@ -326,6 +338,9 @@ class BaseCheckpointSaver(ABC): Returns: RunnableConfig: Updated configuration after storing the checkpoint. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. """ raise NotImplementedError diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index bd5dc6fd1..6d85fd188 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -107,15 +107,16 @@ class MemorySaver(BaseCheckpointSaver): """List checkpoints from the in-memory storage. This method retrieves a list of checkpoint tuples from the in-memory storage based - on the provided config. The checkpoints are ordered by timestamp in insertion order. + on the provided criteria. Args: - config (RunnableConfig): The config to use for listing the checkpoints. - before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None. - limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None. + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): List checkpoints created before this configuration. + limit (Optional[int]): Maximum number of checkpoints to return. Yields: - Iterator[CheckpointTuple]: An iterator of checkpoint tuples. + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage for thread_id in thread_ids: @@ -158,6 +159,7 @@ class MemorySaver(BaseCheckpointSaver): Args: config (RunnableConfig): The config to associate with the checkpoint. checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. Returns: RunnableConfig: The updated config containing the saved checkpoint's timestamp. @@ -191,6 +193,7 @@ class MemorySaver(BaseCheckpointSaver): Args: config (RunnableConfig): The config to associate with the writes. writes (list[tuple[str, Any]]): The writes to save. + task_id (str): Identifier for the task creating the writes. Returns: RunnableConfig: The updated config containing the saved writes' timestamp. @@ -254,6 +257,16 @@ class MemorySaver(BaseCheckpointSaver): checkpoint: Checkpoint, metadata: CheckpointMetadata, ) -> RunnableConfig: + """Asynchronous version of put. + + Args: + config (RunnableConfig): The config to associate with the checkpoint. + checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. + + Returns: + RunnableConfig: The updated config containing the saved checkpoint's timestamp. + """ return await asyncio.get_running_loop().run_in_executor( None, self.put, config, checkpoint, metadata ) @@ -264,6 +277,16 @@ class MemorySaver(BaseCheckpointSaver): writes: List[Tuple[str, Any]], task_id: str, ) -> RunnableConfig: + """Asynchronous version of put_writes. + + This method is an asynchronous wrapper around put_writes that runs the synchronous + method in a separate thread using asyncio. + + Args: + config (RunnableConfig): The config to associate with the writes. + writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair. + task_id (str): Identifier for the task creating the writes. + """ return await asyncio.get_running_loop().run_in_executor( None, self.put_writes, config, writes, task_id ) diff --git a/libs/langgraph/langgraph/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index 6ac9ee593..eee6e05c7 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -309,6 +309,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): Args: config (RunnableConfig): The config to use for listing the checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None. before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None. limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None. @@ -410,6 +411,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): writes: Sequence[Tuple[str, Any]], task_id: str, ) -> None: + """Store intermediate writes linked to a checkpoint. + + This method saves intermediate writes associated with a checkpoint to the SQLite database. + + Args: + 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. + """ with self.lock, self.cursor() as cur: cur.executemany( "INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)", @@ -467,6 +477,17 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): raise NotImplementedError(_AIO_ERROR_MSG) def get_next_version(self, current: Optional[str], channel: BaseChannel) -> str: + """Generate the next version ID for a channel. + + This method creates a new version identifier for a channel based on its current version. + + Args: + current (Optional[str]): The current version identifier of the channel. + channel (BaseChannel): The channel being versioned. + + Returns: + str: The next version identifier, which is guaranteed to be monotonically increasing. + """ if current is None: current_v = 0 else: