Custom loads support in postgres checkpointer (#1930)

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(

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"],

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"