From 29f58fd9e555ab57275e9540ef7b77cf5cea2ed6 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 30 Sep 2024 16:42:23 -0700 Subject: [PATCH] Custom loads support in postgres checkpointer (#1930) --- .../langgraph/store/postgres/aio.py | 32 +++++++++-- .../langgraph/store/postgres/base.py | 53 ++++++++++++++++--- libs/checkpoint-postgres/pyproject.toml | 2 +- 3 files changed, 76 insertions(+), 11 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index 4e59918e0..ef73482ea 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -1,8 +1,18 @@ import asyncio import logging from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Iterable, Sequence, cast +from typing import ( + Any, + AsyncIterator, + Callable, + Iterable, + Optional, + Sequence, + Union, + cast, +) +import orjson from psycopg import AsyncConnection, AsyncCursor from psycopg.errors import UndefinedTable from psycopg.rows import dict_row @@ -19,7 +29,16 @@ logger = logging.getLogger(__name__) class AsyncPostgresStore(BasePostgresStore[AsyncConnection]): - def __init__(self, conn: AsyncConnection[Any]) -> None: + def __init__( + self, + conn: AsyncConnection[Any], + *, + deserializer: Optional[ + Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] + ] = None, + ) -> None: + super().__init__(deserializer=deserializer) + self.conn = conn self.conn = conn self.loop = asyncio.get_running_loop() @@ -87,7 +106,9 @@ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]): for idx, key in items: row = key_to_row.get(key) if row: - results[idx] = _row_to_item(namespace, row) + results[idx] = _row_to_item( + namespace, row, loader=self._deserializer + ) else: results[idx] = None @@ -115,7 +136,10 @@ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]): for cur, idx, op in cursors: rows = cast(list[Row], await cur.fetchall()) - items = [_row_to_item(op.namespace_prefix, row) for row in rows] + items = [ + _row_to_item(op.namespace_prefix, row, loader=self._deserializer) + for row in rows + ] results[idx] = items async def _batch_list_namespaces_ops( diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index a22001355..ba4a39d1f 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -4,7 +4,18 @@ import logging from collections import defaultdict from contextlib import contextmanager from datetime import datetime -from typing import Any, Generic, Iterable, Iterator, Sequence, TypeVar, Union, cast +from typing import ( + Any, + Callable, + Generic, + Iterable, + Iterator, + Optional, + Sequence, + TypeVar, + Union, + cast, +) import orjson from psycopg import BaseConnection, Connection, Cursor @@ -54,6 +65,17 @@ C = TypeVar("C", bound=BaseConnection) class BasePostgresStore(BaseStore, Generic[C]): MIGRATIONS = MIGRATIONS conn: C + __slots__ = ("_deserializer",) + + def __init__( + self, + *, + deserializer: Optional[ + Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] + ] = None, + ) -> None: + super().__init__() + self._deserializer = deserializer def _get_batch_GET_ops_queries( self, @@ -191,7 +213,15 @@ class BasePostgresStore(BaseStore, Generic[C]): class PostgresStore(BasePostgresStore[Connection]): - def __init__(self, conn: Connection[Any]) -> None: + def __init__( + self, + conn: Connection[Any], + *, + deserializer: Optional[ + Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] + ] = None, + ) -> None: + super().__init__(deserializer=deserializer) self.conn = conn def batch(self, ops: Iterable[Op]) -> list[Result]: @@ -246,7 +276,9 @@ class PostgresStore(BasePostgresStore[Connection]): for idx, key in items: row = key_to_row.get(key) if row: - results[idx] = _row_to_item(namespace, row) + results[idx] = _row_to_item( + namespace, row, loader=self._deserializer + ) else: results[idx] = None @@ -274,7 +306,10 @@ class PostgresStore(BasePostgresStore[Connection]): for cur, idx, op in cursors: rows = cast(list[Row], cur.fetchall()) - items = [_row_to_item(op.namespace_prefix, row) for row in rows] + items = [ + _row_to_item(op.namespace_prefix, row, loader=self._deserializer) + for row in rows + ] results[idx] = items def _batch_list_namespaces_ops( @@ -361,11 +396,17 @@ def _namespace_to_ltree(namespace: tuple[str, ...]) -> str: return ".".join(namespace) -def _row_to_item(namespace: tuple[str, ...], row: Row) -> Item: +def _row_to_item( + namespace: tuple[str, ...], + row: Row, + *, + loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None, +) -> Item: """Convert a row from the database into an Item.""" + loader = loader or _json_loads val = row["value"] return Item( - value=val if isinstance(val, dict) else _json_loads(val), + value=val if isinstance(val, dict) else loader(val), key=row["key"], namespace=namespace, created_at=row["created_at"], diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index b65f545e0..e48585900 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-postgres" -version = "1.0.10" +version = "1.0.11" description = "Library with a Postgres implementation of LangGraph checkpoint saver." authors = [] license = "MIT"