merge

libs/checkpoint-postgres/langgraph/store/postgres/aio.py

@@ -17,7 +17,6 @@ from langgraph.store.base import (
     Op,
     PutOp,
     Result,
-    SearchItem,
     SearchOp,
     ensure_embeddings,
 )
@@ -30,6 +29,7 @@ from langgraph.store.postgres.base import (
     _decode_ns_bytes,
     _group_ops,
     _row_to_item,
+    _row_to_search_item,
 )
 
 if TYPE_CHECKING:
@@ -195,11 +195,10 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
             await cur.execute(query, params)
             rows = cast(list[Row], await cur.fetchall())
             items = [
-                _row_to_item(
+                _row_to_search_item(
                     _decode_ns_bytes(row["prefix"]),
                     row,
                     loader=self._deserializer,
-                    cls=SearchItem,
                 )
                 for row in rows
             ]
@@ -258,7 +257,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                     ):
                         yield cur
             else:
-                async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                async with (
+                    self.lock,
+                    conn.transaction(),
+                    conn.cursor(binary=True) as cur,
+                ):
                     yield cur
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:

libs/checkpoint-postgres/langgraph/store/postgres/base.py

@@ -37,6 +37,7 @@ from langgraph.store.base import (
     ListNamespacesOp,
     Op,
     PutOp,
+    ResponseMetadata,
     Result,
     SearchItem,
     SearchOp,
@@ -757,11 +758,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             cur.execute(query, params)
             rows = cast(list[Row], cur.fetchall())
             results[idx] = [
-                _row_to_item(
-                    _decode_ns_bytes(row["prefix"]),
-                    row,
-                    loader=self._deserializer,
-                    cls=SearchItem,
+                _row_to_search_item(
+                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
                 )
                 for row in rows
             ]
@@ -878,6 +876,32 @@ def _row_to_item(
     return cls(**kwargs)
 
 
+def _row_to_search_item(
+    namespace: tuple[str, ...],
+    row: Row,
+    *,
+    loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
+) -> SearchItem:
+    """Convert a row from the database into an Item."""
+    loader = loader or _json_loads
+    val = row["value"]
+    response_metadata: Optional[ResponseMetadata] = (
+        {
+            "score": float(row["score"]),
+        }
+        if row.get("score") is not None
+        else None
+    )
+    return SearchItem(
+        value=val if isinstance(val, dict) else loader(val),
+        key=row["key"],
+        namespace=namespace,
+        created_at=row["created_at"],
+        updated_at=row["updated_at"],
+        response_metadata=response_metadata,
+    )
+
+
 def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
     grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
     tot = 0

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.3"
+version = "2.0.4"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"