Skip to content

Commit

Permalink
checkpoint postgres: allow passing custom serde (#1699)
Browse files Browse the repository at this point in the history
* Use correct serde in postgres checkpoint

* Revert metadata serde

* Revert

---------

Co-authored-by: Vadym Barda <vadym@langchain.dev>
  • Loading branch information
vermapratyush and vbarda committed Sep 13, 2024
1 parent 66fc7c9 commit c3b4acf
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,11 @@ def __init__(
@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.
Expand All @@ -73,9 +77,9 @@ async def from_conn_string(
) 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.
Expand Down

0 comments on commit c3b4acf

Please sign in to comment.