Add in-mem vector search (#2547)

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -378,15 +378,19 @@ class PostgresSaver(BasePostgresSaver):
                 # a connection not in pipeline mode can only be used by one
                 # thread/coroutine at a time, so we acquire a lock
                 if self.supports_pipeline:
-                    with self.lock, conn.pipeline(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    with (
+                        self.lock,
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
                 else:
                     # Use connection's transaction context manager when pipeline mode not supported
-                    with self.lock, conn.transaction(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    with (
+                        self.lock,
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
             else:
                 with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:

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

@@ -338,20 +338,25 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 # a connection not in pipeline mode can only be used by one
                 # thread/coroutine at a time, so we acquire a lock
                 if self.supports_pipeline:
-                    async with self.lock, conn.pipeline(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    async with (
+                        self.lock,
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
                 else:
                     # Use connection's transaction context manager when pipeline mode not supported
-                    async with self.lock, conn.transaction(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    async with (
+                        self.lock,
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
             else:
-                async with self.lock, conn.cursor(
-                    binary=True, row_factory=dict_row
-                ) as cur:
+                async with (
+                    self.lock,
+                    conn.cursor(binary=True, row_factory=dict_row) as cur,
+                ):
                     yield cur
 
     def list(

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

@@ -28,6 +28,7 @@ from langgraph.store.postgres.base import (
     _decode_ns_bytes,
     _group_ops,
     _row_to_item,
+    _row_to_search_item,
 )
 
 logger = logging.getLogger(__name__)
@@ -146,7 +147,7 @@ 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
                 )
                 for row in rows
@@ -195,9 +196,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                 async with self.lock, conn.pipeline(), conn.cursor(binary=True) as cur:
                     yield cur
             else:
-                async with self.lock, conn.transaction(), conn.cursor(
-                    binary=True
-                ) as cur:
+                async with (
+                    self.lock,
+                    conn.transaction(),
+                    conn.cursor(binary=True) as cur,
+                ):
                     yield cur
         else:
             async with conn.cursor(binary=True) as cur:

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

@@ -36,6 +36,7 @@ from langgraph.store.base import (
     Op,
     PutOp,
     Result,
+    SearchItem,
     SearchOp,
 )
 
@@ -344,14 +345,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                 # a connection not in pipeline mode can only be used by one
                 # thread/coroutine at a time, so we acquire a lock
                 if self.supports_pipeline:
-                    with self.lock, conn.pipeline(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    with (
+                        self.lock,
+                        conn.pipeline(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
                 else:
-                    with self.lock, conn.transaction(), conn.cursor(
-                        binary=True, row_factory=dict_row
-                    ) as cur:
+                    with (
+                        self.lock,
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
                         yield cur
             else:
                 with conn.cursor(binary=True, row_factory=dict_row) as cur:
@@ -430,7 +435,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             cur.execute(query, params)
             rows = cast(list[Row], cur.fetchall())
             results[idx] = [
-                _row_to_item(
+                _row_to_search_item(
                     _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
                 )
                 for row in rows
@@ -517,6 +522,32 @@ def _row_to_item(
     )
 
 
+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"]
+    score = row.get("score")
+    if score is not None:
+        try:
+            score = float(score)  # type: ignore[arg-type]
+        except ValueError:
+            logger.warning("Invalid score: %s", score)
+            score = 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"],
+        score=score,
+    )
+
+
 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