From c3b4acf3f18df1766ffb1c40e384acbbfb1c7122 Mon Sep 17 00:00:00 2001 From: Pratyush Verma Date: Fri, 13 Sep 2024 05:09:23 +0100 Subject: [PATCH] checkpoint postgres: allow passing custom serde (#1699) * Use correct serde in postgres checkpoint * Revert metadata serde * Revert --------- Co-authored-by: Vadym Barda --- .../langgraph/checkpoint/postgres/aio.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 9c246ef74..9aa08ad0f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -57,7 +57,11 @@ class AsyncPostgresSaver(BasePostgresSaver): @classmethod @asynccontextmanager async def from_conn_string( - cls, conn_string: str, *, pipeline: bool = False + cls, + conn_string: str, + *, + pipeline: bool = False, + serde: Optional[SerializerProtocol] = None, ) -> AsyncIterator["AsyncPostgresSaver"]: """Create a new PostgresSaver instance from a connection string. @@ -73,9 +77,9 @@ class AsyncPostgresSaver(BasePostgresSaver): ) as conn: if pipeline: async with conn.pipeline() as pipe: - yield AsyncPostgresSaver(conn, pipe) + yield AsyncPostgresSaver(conn=conn, pipe=pipe, serde=serde) else: - yield AsyncPostgresSaver(conn) + yield AsyncPostgresSaver(conn=conn, serde=serde) async def setup(self) -> None: """Set up the checkpoint database asynchronously.