diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 65f92b8d7..be052ec86 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -66,9 +66,9 @@ async def from_conn_string( ) as conn: if pipeline: async with conn.pipeline() as pipe: - yield AsyncPostgresSaver(conn=conn, pipe=pipe, serde=serde) + yield cls(conn=conn, pipe=pipe, serde=serde) else: - yield AsyncPostgresSaver(conn=conn, serde=serde) + yield cls(conn=conn, serde=serde) async def setup(self) -> None: """Set up the checkpoint database asynchronously. @@ -143,15 +143,17 @@ async def alist( value["pending_sends"], ), self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], + ( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - } - if value["parent_checkpoint_id"] - else None, + if value["parent_checkpoint_id"] + else None + ), await asyncio.to_thread(self._load_writes, value["pending_writes"]), ) @@ -202,15 +204,17 @@ async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: value["pending_sends"], ), self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["parent_checkpoint_id"], + ( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": value["parent_checkpoint_id"], + } } - } - if value["parent_checkpoint_id"] - else None, + if value["parent_checkpoint_id"] + else None + ), await asyncio.to_thread(self._load_writes, value["pending_writes"]), )