diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 63a2038a3..a931181e1 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,6 +1,6 @@ import asyncio from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Optional, Union +from typing import Any, AsyncIterator, Iterator, List, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline @@ -51,6 +51,7 @@ class AsyncPostgresSaver(BasePostgresSaver): self.conn = conn self.pipe = pipe self.lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() @classmethod @asynccontextmanager @@ -329,3 +330,96 @@ class AsyncPostgresSaver(BasePostgresSaver): binary=True, row_factory=dict_row ) as cur: yield cur + + def list( + self, + config: RunnableConfig | None, + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[CheckpointTuple]: + """List checkpoints from the database. + + This method retrieves a list of checkpoint tuples from the Postgres database based + on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). + + Args: + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + limit (Optional[int]): Maximum number of checkpoints to return. + + Yields: + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. + """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), self.loop + ).result() + except StopAsyncIteration: + break + + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + """Get a checkpoint tuple from the database. + + This method retrieves a checkpoint tuple from the Postgres database based on the + provided config. If the config contains a "checkpoint_id" key, the checkpoint with + the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint + for the given thread ID is retrieved. + + Args: + config (RunnableConfig): The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() + + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + """Save a checkpoint to the database. + + This method saves a checkpoint to the Postgres database. The checkpoint is associated + with the provided config and its parent config (if any). + + 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. + new_versions (ChannelVersions): New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, + config: RunnableConfig, + writes: List[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 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. + """ + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id), self.loop + ).result() diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index dc68a0403..7cdc7c8fb 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -1,11 +1,11 @@ import asyncio -import functools from contextlib import asynccontextmanager from typing import ( Any, AsyncIterator, Dict, Iterator, + List, Optional, Sequence, Tuple, @@ -31,20 +31,6 @@ from langgraph.checkpoint.sqlite.utils import search_where T = TypeVar("T", bound=callable) -def not_implemented_sync_method(func: T) -> T: - @functools.wraps(func) - def wrapper(*args, **kwargs): - raise NotImplementedError( - "The AsyncSqliteSaver does not support synchronous methods. " - "Consider using the SqliteSaver instead.\n" - "from langgraph.checkpoint.sqlite import SqliteSaver\n" - "See https://langchain-ai.github.io/langgraph/reference/checkpoints/langgraph.checkpoint.sqlite.SqliteSaver " - "for more information." - ) - - return wrapper - - class AsyncSqliteSaver(BaseCheckpointSaver): """An asynchronous checkpoint saver that stores checkpoints in a SQLite database. @@ -132,6 +118,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): self.jsonplus_serde = JsonPlusSerializer() self.conn = conn self.lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() self.is_setup = False @classmethod @@ -150,16 +137,24 @@ class AsyncSqliteSaver(BaseCheckpointSaver): async with aiosqlite.connect(conn_string) as conn: yield AsyncSqliteSaver(conn) - @not_implemented_sync_method def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. - Note: - This method is not implemented for the AsyncSqliteSaver. Use `aget` instead. - Or consider using the [SqliteSaver][sqlitesaver] checkpointer. - """ + This method retrieves a checkpoint tuple from the SQLite database based on the + provided config. If the config contains a "checkpoint_id" key, the checkpoint with + the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint + for the given thread ID is retrieved. + + Args: + config (RunnableConfig): The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() - @not_implemented_sync_method def list( self, config: Optional[RunnableConfig], @@ -168,21 +163,60 @@ class AsyncSqliteSaver(BaseCheckpointSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. + """List checkpoints from the database asynchronously. - Note: - This method is not implemented for the AsyncSqliteSaver. Use `alist` instead. - Or consider using the [SqliteSaver][sqlitesaver] checkpointer. + This method retrieves a list of checkpoint tuples from the SQLite database based + on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). + + Args: + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + limit (Optional[int]): Maximum number of checkpoints to return. + + Yields: + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), self.loop + ).result() + except StopAsyncIteration: + break - @not_implemented_sync_method def put( self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, + new_versions: ChannelVersions, ) -> RunnableConfig: - """Save a checkpoint to the database. FOO""" + """Save a checkpoint to the database. + + This method saves a checkpoint to the SQLite database. The checkpoint is associated + with the provided config and its parent config (if any). + + 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. + new_versions (ChannelVersions): New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str + ) -> None: + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id), self.loop + ).result() async def setup(self) -> None: """Set up the checkpoint database asynchronously. diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d3a6c422f..7901ad2bc 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8771,8 +8771,10 @@ async def test_weather_subgraph( else: return "normal_llm_node" - async def weather_graph(state: RouterState): - return await subgraph.ainvoke(state) + def weather_graph(state: RouterState): + # this tests that all async checkpointers tested also implement sync methods + # as the subgraph called with sync invoke will use sync checkpointer methods + return subgraph.invoke(state) graph = StateGraph(RouterState) graph.add_node(router_node) @@ -8783,6 +8785,9 @@ async def test_weather_subgraph( graph.add_edge("normal_llm_node", END) graph.add_edge("weather_graph", END) + def get_first_in_list(): + return [*graph.get_state_history(config, limit=1)][0] + async with awith_checkpointer(checkpointer_name) as checkpointer: graph = graph.compile(checkpointer=checkpointer) @@ -8846,6 +8851,8 @@ async def test_weather_subgraph( ), ), ) + # confirm that list() delegates to alist() correctly + assert await asyncio.to_thread(get_first_in_list) == state # update await graph.aupdate_state(state.tasks[0].state, {"city": "la"})