Skip to content

Commit

Permalink
Use correct serde in postgres checkpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
vermapratyush committed Sep 12, 2024
1 parent 889cc9b commit 7be933f
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 4 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
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,7 @@ def _dump_writes(
]

def _load_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata))
return self.serde.loads(self.serde.dumps(metadata))

def _dump_metadata(self, metadata) -> str:
serialized_metadata_type, serialized_metadata = self.jsonplus_serde.dumps_typed(
Expand Down

0 comments on commit 7be933f

Please sign in to comment.