Add optoin to not index

libs/checkpoint-postgres/Makefile

@@ -5,7 +5,11 @@
 ######################
 
 start-postgres:
-	POSTGRES_VERSION=${POSTGRES_VERSION:-16} docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait
+	POSTGRES_VERSION=${POSTGRES_VERSION:-16} docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait || ( \
+		echo "Failed to start PostgreSQL, printing logs..."; \
+		docker compose -f tests/compose-postgres.yml logs; \
+		exit 1 \
+	)
 
 stop-postgres:
 	docker compose -f tests/compose-postgres.yml down

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

@@ -1,6 +1,7 @@
 import threading
+from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
-from typing import Any, Iterator, Optional, Sequence
+from typing import Any, Optional
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import Capabilities, Connection, Cursor, Pipeline
@@ -378,15 +379,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/_ainternal.py

@@ -1,7 +1,8 @@
 """Shared async utility functions for the Postgres checkpoint & storage classes."""
 
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from typing import AsyncIterator, Union
+from typing import Union
 
 from psycopg import AsyncConnection
 from psycopg.rows import DictRow

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

@@ -1,7 +1,8 @@
 """Shared utility functions for the Postgres checkpoint & storage classes."""
 
+from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Iterator, Union
+from typing import Union
 
 from psycopg import Connection
 from psycopg.rows import DictRow

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

@@ -1,6 +1,7 @@
 import asyncio
+from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Iterator, Optional, Sequence
+from typing import Any, Optional
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
@@ -338,20 +339,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(
@@ -380,7 +386,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         while True:
             try:
                 yield asyncio.run_coroutine_threadsafe(
-                    anext(aiter_),
+                    anext(aiter_),  # noqa: F821
                     self.loop,
                 ).result()
             except StopAsyncIteration:

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

@@ -1,5 +1,6 @@
 import random
-from typing import Any, List, Optional, Sequence, Tuple, cast
+from collections.abc import Sequence
+from typing import Any, Optional, cast
 
 from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
@@ -249,7 +250,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
         config: Optional[RunnableConfig],
         filter: MetadataInput,
         before: Optional[RunnableConfig] = None,
-    ) -> Tuple[str, List[Any]]:
+    ) -> tuple[str, list[Any]]:
         """Return WHERE clause predicates for alist() given config, filter, before.
 
         This method returns a tuple of a string and a tuple of values. The string

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

@@ -1,16 +1,8 @@
 import asyncio
 import logging
+from collections.abc import AsyncIterator, Iterable, Sequence
 from contextlib import asynccontextmanager
-from typing import (
-    Any,
-    AsyncIterator,
-    Callable,
-    Iterable,
-    Optional,
-    Sequence,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
 
 import orjson
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
@@ -19,22 +11,41 @@ from psycopg.rows import DictRow, dict_row
 from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.postgres import _ainternal
-from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
+from langgraph.store.base import (
+    GetOp,
+    ListNamespacesOp,
+    Op,
+    PutOp,
+    Result,
+    SearchItem,
+    SearchOp,
+    ensure_embeddings,
+)
 from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.postgres.base import (
     BasePostgresStore,
     PoolConfig,
+    PostgresEmbeddingConfig,
     Row,
     _decode_ns_bytes,
     _group_ops,
     _row_to_item,
 )
 
+if TYPE_CHECKING:
+    from langchain_core.embeddings import Embeddings
+
 logger = logging.getLogger(__name__)
 
 
 class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
-    __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
+    __slots__ = (
+        "_deserializer",
+        "pipe",
+        "lock",
+        "supports_pipeline",
+        "embedding_config",
+    )
 
     def __init__(
         self,
@@ -44,6 +55,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         deserializer: Optional[
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
+        embedding: Optional[PostgresEmbeddingConfig] = None,
     ) -> None:
         if isinstance(conn, AsyncConnectionPool) and pipe is not None:
             raise ValueError(
@@ -56,6 +68,14 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         self.lock = asyncio.Lock()
         self.loop = asyncio.get_running_loop()
         self.supports_pipeline = Capabilities().has_pipeline()
+        self.embedding_config = embedding
+        if self.embedding_config:
+            self.embeddings: Optional[Embeddings] = ensure_embeddings(
+                self.embedding_config.get("embed"),
+                aembed=self.embedding_config.get("aembed"),
+            )
+        else:
+            self.embeddings = None
 
     async def abatch(self, ops: Iterable[Op]) -> list[Result]:
         grouped_ops, num_ops = _group_ops(ops)
@@ -76,7 +96,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         results: list[Result],
         conn: AsyncConnection[DictRow],
     ) -> None:
-        async with self._cursor(conn, pipeline=True) as cur:
+        async with self._cursor(pipeline=True) as cur:
             if GetOp in grouped_ops:
                 await self._batch_get_ops(
                     cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]),
@@ -131,7 +151,28 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         put_ops: Sequence[tuple[int, PutOp]],
         cur: AsyncCursor[DictRow],
     ) -> None:
-        queries = self._get_batch_PUT_queries(put_ops)
+        queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
+        if embedding_request:
+            if self.embeddings is None:
+                # Should not get here since the embedding config is required
+                # to return an embedding_request above
+                raise ValueError(
+                    "Embedding configuration is required for vector operations "
+                    f"(for semantic search). "
+                    f"Please provide an EmbeddingConfig when initializing the {self.__class__.__name__}."
+                )
+            query, txt_params = embedding_request
+            # Update the params to replace the raw text with the vectors
+            vectors = await self.embeddings.aembed_documents(
+                [param[-1] for param in txt_params]
+            )
+            queries.extend(
+                [
+                    (query, (ns, key, value, vector))
+                    for (ns, key, value, _), vector in zip(txt_params, vectors)
+                ]
+            )
+
         for query, params in queries:
             await cur.execute(query, params)
 
@@ -141,13 +182,24 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         results: list[Result],
         cur: AsyncCursor[DictRow],
     ) -> None:
-        queries = self._get_batch_search_queries(search_ops)
-        for (query, params), (idx, _) in zip(queries, search_ops):
+        queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
+
+        if embedding_requests and self.embeddings:
+            embeddings = await self.embeddings.aembed_documents(
+                [query for _, query in embedding_requests]
+            )
+            for (idx, _), embedding in zip(embedding_requests, embeddings):
+                queries[idx][1][0] = embedding
+
+        for (idx, _), (query, params) in zip(search_ops, queries):
             await cur.execute(query, params)
             rows = cast(list[Row], await cur.fetchall())
             items = [
                 _row_to_item(
-                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                    _decode_ns_bytes(row["prefix"]),
+                    row,
+                    loader=self._deserializer,
+                    cls=SearchItem,
                 )
                 for row in rows
             ]
@@ -168,40 +220,46 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
 
     @asynccontextmanager
     async def _cursor(
-        self, conn: AsyncConnection[DictRow], *, pipeline: bool = False
-    ) -> AsyncIterator[AsyncCursor[Any]]:
+        self, *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[DictRow]]:
         """Create a database cursor as a context manager.
 
         Args:
-            conn: The database connection to use
             pipeline: whether to use pipeline for the DB operations inside the context manager.
                 Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline.
                 If pipeline mode is not supported, will fall back to using transaction context manager.
         """
-        if self.pipe:
-            # a connection in pipeline mode can be used concurrently
-            # in multiple threads/coroutines, but only one cursor can be
-            # used at a time
-            async with conn.cursor(binary=True) as cur:
+        async with _ainternal.get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
                 try:
-                    yield cur
+                    async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
                 finally:
                     if pipeline:
                         await self.pipe.sync()
-        elif pipeline:
-            # 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) as cur:
-                    yield cur
+            elif pipeline:
+                # 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,
+                    ):
+                        yield cur
+                else:
+                    async with (
+                        self.lock,
+                        conn.transaction(),
+                        conn.cursor(binary=True, row_factory=dict_row) as cur,
+                    ):
+                        yield cur
             else:
-                async with self.lock, conn.transaction(), conn.cursor(
-                    binary=True
-                ) as cur:
+                async with conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
-        else:
-            async with conn.cursor(binary=True) as cur:
-                yield cur
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:
         return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
@@ -214,6 +272,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         *,
         pipeline: bool = False,
         pool_config: Optional[PoolConfig] = None,
+        embedding: Optional[PostgresEmbeddingConfig] = None,
     ) -> AsyncIterator["AsyncPostgresStore"]:
         """Create a new AsyncPostgresStore instance from a connection string.
 
@@ -223,6 +282,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
             pool_config (Optional[PoolConfig]): Configuration for the connection pool.
                 If provided, will create a connection pool and use it instead of a single connection.
                 This overrides the `pipeline` argument.
+            embedding (Optional[PostgresEmbeddingConfig]): The embedding config.
 
         Returns:
             AsyncPostgresStore: A new AsyncPostgresStore instance.
@@ -244,16 +304,16 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                     **cast(dict, pc),
                 ),
             ) as pool:
-                yield cls(conn=pool)
+                yield cls(conn=pool, embedding=embedding)
         else:
             async with await AsyncConnection.connect(
                 conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
             ) as conn:
                 if pipeline:
                     async with conn.pipeline() as pipe:
-                        yield cls(conn=conn, pipe=pipe)
+                        yield cls(conn=conn, pipe=pipe, embedding=embedding)
                 else:
-                    yield cls(conn=conn)
+                    yield cls(conn=conn, embedding=embedding)
 
     async def setup(self) -> None:
         """Set up the store database asynchronously.
@@ -262,33 +322,45 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         already exist and runs database migrations. It MUST be called directly by the user
         the first time the store is used.
         """
-        async with _ainternal.get_connection(self.conn) as conn:
-            async with conn.cursor() as cur:
-                try:
-                    await cur.execute(
-                        "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
-                    )
-                    row = cast(dict, await cur.fetchone())
-                    if row is None:
-                        version = -1
-                    else:
-                        version = row["v"]
-                except UndefinedTable:
+        async with self._cursor() as cur:
+            try:
+                await cur.execute(
+                    "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
+                )
+                row = await cur.fetchone()
+                if row is None:
                     version = -1
-                    # Create store_migrations table if it doesn't exist
-                    await cur.execute(
-                        """
-                        CREATE TABLE IF NOT EXISTS store_migrations (
-                            v INTEGER PRIMARY KEY
-                        )
-                        """
-                    )
-                for v, migration in enumerate(
-                    self.MIGRATIONS[version + 1 :], start=version + 1
-                ):
-                    await cur.execute(migration)
-                    await cur.execute(
-                        "INSERT INTO store_migrations (v) VALUES (%s)", (v,)
+                else:
+                    version = row["v"]
+            except UndefinedTable:
+                version = -1
+                await cur.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS store_migrations (
+                        v INTEGER PRIMARY KEY
                     )
+                """
+                )
+
+            for v, migration in enumerate(
+                self.MIGRATIONS[version + 1 :], start=version + 1
+            ):
+                if isinstance(migration, str):
+                    sql = migration
+                else:
+                    if migration.condition and not migration.condition(self):
+                        continue
+
+                    sql = migration.sql
+                    if migration.params:
+                        params = {
+                            k: v(self) if v is not None and callable(v) else v
+                            for k, v in migration.params.items()
+                        }
+                        sql = sql % params
+
+                await cur.execute(sql)
+                await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+
             if self.pipe:
                 await self.pipe.sync()

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

@@ -3,16 +3,17 @@ import json
 import logging
 import threading
 from collections import defaultdict
+from collections.abc import Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from datetime import datetime
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Generic,
-    Iterable,
-    Iterator,
+    Literal,
+    NamedTuple,
     Optional,
-    Sequence,
     TypeVar,
     Union,
     cast,
@@ -30,19 +31,80 @@ from langgraph.checkpoint.postgres import _ainternal as _ainternal
 from langgraph.checkpoint.postgres import _internal as _pg_internal
 from langgraph.store.base import (
     BaseStore,
+    EmbeddingConfig,
     GetOp,
     Item,
     ListNamespacesOp,
     Op,
     PutOp,
     Result,
+    SearchItem,
     SearchOp,
+    ensure_embeddings,
 )
 
+if TYPE_CHECKING:
+    from langchain_core.embeddings import Embeddings
+
 logger = logging.getLogger(__name__)
 
 
-MIGRATIONS = [
+class Migration(NamedTuple):
+    """A database migration with optional conditions and parameters."""
+
+    sql: str
+    condition: Optional[Callable[[Any], bool]] = None
+    params: Optional[dict[str, Any]] = None
+
+
+def _embedding_requested(store: Any) -> bool:
+    """Check if vector operations are available in the database."""
+    return bool(store.embedding_config)
+
+
+def _get_vector_type_ops(store: Any) -> str:
+    """Get the vector type operator class based on config."""
+    if not store.embedding_config:
+        return "vector_cosine_ops"
+
+    config = cast(PostgresEmbeddingConfig, store.embedding_config)
+    index_config = config.get(
+        "index_config", BasePostgresStore._get_default_index_config()
+    )
+    vector_type = index_config.get("vector_type", "vector")
+    distance_type = config.get("distance_type", "cosine")
+
+    # For regular vectors
+    type_prefix = {"vector": "vector", "halfvec": "halfvec"}[vector_type]
+
+    if distance_type not in ("l2", "inner_product", "cosine"):
+        raise ValueError(
+            f"Vector type {vector_type} only supports 'l2', 'inner_product', or 'cosine' distance, got {distance_type}"
+        )
+
+    distance_suffix = {
+        "l2": "l2_ops",
+        "inner_product": "ip_ops",
+        "cosine": "cosine_ops",
+    }[distance_type]
+
+    return f"{type_prefix}_{distance_suffix}"
+
+
+def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]:
+    """Get the index type and configuration based on config."""
+    if not store.embedding_config:
+        return "hnsw", {}
+
+    config = cast(PostgresEmbeddingConfig, store.embedding_config)
+    default_config = BasePostgresStore._get_default_index_config()
+    index_config = config.get("index_config", default_config).copy()
+    kind = index_config.pop("kind", "hnsw")
+    index_config.pop("vector_type", None)
+    return kind, index_config
+
+
+MIGRATIONS: Sequence[Union[str, Migration]] = [
     """
 CREATE TABLE IF NOT EXISTS store (
     -- 'prefix' represents the doc's 'namespace'
@@ -58,8 +120,56 @@ CREATE TABLE IF NOT EXISTS store (
 -- For faster lookups by prefix
 CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
 """,
+    Migration(
+        """
+CREATE EXTENSION IF NOT EXISTS vector;
+""",
+        condition=_embedding_requested,
+    ),
+    Migration(
+        """
+CREATE TABLE IF NOT EXISTS store_vectors (
+    prefix text NOT NULL,
+    key text NOT NULL,
+    field_name text NOT NULL,
+    embedding %(vector_type)s(%(dims)s),
+    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
+    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (prefix, key, field_name),
+    FOREIGN KEY (prefix, key) REFERENCES store(prefix, key) ON DELETE CASCADE
+);
+""",
+        condition=_embedding_requested,
+        params={
+            "dims": lambda store: store.embedding_config["dims"],
+            "vector_type": lambda store: (
+                cast(PostgresEmbeddingConfig, store.embedding_config)
+                .get("index_config", {})
+                .get("vector_type", "vector")
+            ),
+        },
+    ),
+    Migration(
+        """
+CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
+    USING %(index_type)s (embedding %(ops)s)%(index_params)s;
+""",
+        condition=_embedding_requested,
+        params={
+            "index_type": lambda store: _get_index_params(store)[0],
+            "ops": lambda store: _get_vector_type_ops(store),
+            "index_params": lambda store: (
+                " WITH ("
+                + ", ".join(f"{k}={v}" for k, v in _get_index_params(store)[1].items())
+                + ")"
+                if _get_index_params(store)[1]
+                else ""
+            ),
+        },
+    ),
 ]
 
+
 C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
 
 
@@ -88,10 +198,76 @@ class PoolConfig(TypedDict, total=False):
     """
 
 
+class IndexConfig(TypedDict, total=False):
+    """Configuration for vector index in PostgreSQL store."""
+
+    kind: Literal["hnsw", "ivfflat"]
+    """Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat."""
+    vector_type: Literal["vector", "halfvec"]
+    """Type of vector storage to use.
+    Options:
+    - 'vector': Regular vectors (default)
+    - 'halfvec': Half-precision vectors for reduced memory usage
+    """
+
+
+class HNSWConfig(IndexConfig, total=False):
+    """Configuration for HNSW (Hierarchical Navigable Small World) index."""
+
+    kind: Literal["hnsw"]  # type: ignore[misc]
+    m: int
+    """Maximum number of connections per layer. Default is 16."""
+    ef_construction: int
+    """Size of dynamic candidate list for index construction. Default is 64."""
+
+
+class IVFFlatConfig(IndexConfig, total=False):
+    """IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
+
+    Three keys to achieving good recall are:
+    1. Create the index after the table has some data
+    2. Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for over 1M rows
+    3. When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) - a good place to start is sqrt(lists)
+    """
+
+    kind: Literal["ivfflat"]  # type: ignore[misc]
+    nlist: int
+    """Number of inverted lists (clusters) for IVF index.
+    
+    Determines the number of clusters used in the index structure.
+    Higher values can improve search speed but increase index size and build time.
+    Typically set to the square root of the number of vectors in the index.
+    """
+
+
+class PostgresEmbeddingConfig(EmbeddingConfig, total=False):
+    """Configuration for vector embeddings in PostgreSQL store with pgvector-specific options.
+
+    Extends EmbeddingConfig with additional configuration for pgvector index and vector types.
+    """
+
+    index_config: Union[HNSWConfig, IVFFlatConfig]
+    """Specific configuration for the chosen index type (HNSW or IVF Flat)."""
+    distance_type: Literal["l2", "inner_product", "cosine"]
+    """Distance metric to use for vector similarity search:
+    - 'l2': Euclidean distance
+    - 'inner_product': Dot product
+    - 'cosine': Cosine similarity
+    """
+
+
 class BasePostgresStore(Generic[C]):
     MIGRATIONS = MIGRATIONS
     conn: C
     _deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]]
+    embedding_config: Optional[PostgresEmbeddingConfig]
+
+    @staticmethod
+    def _get_default_index_config() -> IndexConfig:
+        return HNSWConfig(
+            kind="hnsw",
+            vector_type="vector",
+        )
 
     def _get_batch_GET_ops_queries(
         self,
@@ -113,10 +289,13 @@ class BasePostgresStore(Generic[C]):
             results.append((query, params, namespace, items))
         return results
 
-    def _get_batch_PUT_queries(
+    def _prepare_batch_PUT_queries(
         self,
         put_ops: Sequence[tuple[int, PutOp]],
-    ) -> list[tuple[str, Sequence]]:
+    ) -> tuple[
+        list[tuple[str, Sequence]],
+        Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
+    ]:
         # Last-write wins
         dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
         for _, op in put_ops:
@@ -143,60 +322,145 @@ class BasePostgresStore(Generic[C]):
                 )
                 params = (_namespace_to_text(namespace), *keys)
                 queries.append((query, params))
+        embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = (
+            None
+        )
         if inserts:
             values = []
             insertion_params = []
+            vector_values = []
+            embedding_request_params = []
+
+            # First handle main store insertions
             for op in inserts:
                 values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
                 insertion_params.extend(
                     [
                         _namespace_to_text(op.namespace),
                         op.key,
-                        Jsonb(op.value),
+                        Jsonb(cast(dict, op.value).copy()),
                     ]
                 )
+
+            # Then handle embeddings if configured
+            if self.embedding_config:
+                text_fields = self.embedding_config.get("text_fields", ["__root__"])
+                if isinstance(text_fields, str):
+                    text_fields = [text_fields]
+                elif text_fields is None:
+                    text_fields = ["__root__"]
+                for op in inserts:
+                    if op.index is False:
+                        continue
+                    value = op.value
+                    ns = _namespace_to_text(op.namespace)
+                    k = op.key
+
+                    for field in text_fields:
+                        for text in _extract_text_by_path(value, field):
+                            vector_values.append(
+                                "(%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
+                            )
+                            embedding_request_params.append((ns, k, field, text))
+
             values_str = ",".join(values)
             query = f"""
                 INSERT INTO store (prefix, key, value, created_at, updated_at)
                 VALUES {values_str}
                 ON CONFLICT (prefix, key) DO UPDATE
-                SET value = EXCLUDED.value, updated_at = CURRENT_TIMESTAMP
+                SET value = EXCLUDED.value,
+                    updated_at = CURRENT_TIMESTAMP
             """
             queries.append((query, insertion_params))
 
-        return queries
+            if vector_values:
+                values_str = ",".join(vector_values)
+                query = f"""
+                    INSERT INTO store_vectors (prefix, key, field_name, embedding, created_at, updated_at)
+                    VALUES {values_str}
+                    ON CONFLICT (prefix, key, field_name) DO UPDATE
+                    SET embedding = EXCLUDED.embedding,
+                        updated_at = CURRENT_TIMESTAMP
+                """
+                embedding_request = (query, embedding_request_params)
 
-    def _get_batch_search_queries(
+        return queries, embedding_request
+
+    def _prepare_batch_search_queries(
         self,
         search_ops: Sequence[tuple[int, SearchOp]],
-    ) -> list[tuple[str, Sequence]]:
-        queries: list[tuple[str, Sequence]] = []
-        for _, op in search_ops:
-            query = """
+    ) -> tuple[
+        list[tuple[str, list[Union[None, str, list[float]]]]],  # queries, params
+        list[tuple[int, str]],  # idx, query_text pairs to embed
+    ]:
+        queries = []
+        embedding_requests = []
+
+        for idx, (_, op) in enumerate(search_ops):
+            base_query = """
                 SELECT prefix, key, value, created_at, updated_at
                 FROM store
                 WHERE prefix LIKE %s
             """
             params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
+            needs_vector_search = False
+
+            if op.query and self.embedding_config:
+                needs_vector_search = True
+                embedding_requests.append((idx, op.query))
+
+                _, score_expr = _get_distance_operator(self)
+                vector_type = (
+                    cast(PostgresEmbeddingConfig, self.embedding_config)
+                    .get("index_config", self._get_default_index_config())
+                    .get("vector_type", "vector")
+                )
+
+                # For hamming distance, we need the vector dimension for normalization
+                if (
+                    vector_type == "bit"
+                    and self.embedding_config.get("distance_type") == "hamming"
+                ):
+                    score_expr = score_expr % ("%s", self.embedding_config["dims"])
+                else:
+                    score_expr = score_expr % ("%s", vector_type)
+
+                base_query = f"""
+                    SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at,
+                           {score_expr} as score
+                    FROM store s
+                    JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
+                    WHERE s.prefix LIKE %s
+                """
+                params = [None, f"{_namespace_to_text(op.namespace_prefix)}%"]
 
             if op.filter:
                 filter_conditions = []
                 for key, value in op.filter.items():
-                    if isinstance(value, list):
-                        filter_conditions.append("value->%s @> %s::jsonb")
-                        params.extend([key, json.dumps(value)])
+                    if isinstance(value, dict):
+                        for op_name, val in value.items():
+                            condition, filter_params = self._get_filter_condition(
+                                key, op_name, val
+                            )
+                            filter_conditions.append(condition)
+                            params.extend(filter_params)
                     else:
                         filter_conditions.append("value->%s = %s::jsonb")
                         params.extend([key, json.dumps(value)])
-                query += " AND " + " AND ".join(filter_conditions)
 
-            # Note: we will need to not do this if sim/keyword search
-            # is used
-            query += " ORDER BY updated_at DESC LIMIT %s OFFSET %s"
+                if filter_conditions:
+                    base_query += " AND " + " AND ".join(filter_conditions)
+
+            order_by = (
+                "ORDER BY score DESC"
+                if needs_vector_search
+                else "ORDER BY updated_at DESC"
+            )
+            base_query += f" {order_by} LIMIT %s OFFSET %s"
             params.extend([op.limit, op.offset])
+            queries.append((base_query, params))
 
-            queries.append((query, params))
-        return queries
+        return queries, embedding_requests
 
     def _get_batch_list_namespaces_queries(
         self,
@@ -248,10 +512,27 @@ class BasePostgresStore(Generic[C]):
 
             query += " ORDER BY truncated_prefix LIMIT %s OFFSET %s"
             params.extend([op.limit, op.offset])
-            queries.append((query, params))
+            queries.append((query, tuple(params)))
 
         return queries
 
+    def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
+        """Helper to generate filter conditions."""
+        if op == "$eq":
+            return "value->%s = %s::jsonb", [key, json.dumps(value)]
+        elif op == "$gt":
+            return "value->>%s > %s", [key, str(value)]
+        elif op == "$gte":
+            return "value->>%s >= %s", [key, str(value)]
+        elif op == "$lt":
+            return "value->>%s < %s", [key, str(value)]
+        elif op == "$lte":
+            return "value->>%s <= %s", [key, str(value)]
+        elif op == "$ne":
+            return "value->%s != %s::jsonb", [key, json.dumps(value)]
+        else:
+            raise ValueError(f"Unsupported operator: {op}")
+
 
 class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
     __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
@@ -264,6 +545,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         deserializer: Optional[
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
+        embedding: Optional[PostgresEmbeddingConfig] = None,
     ) -> None:
         super().__init__()
         self._deserializer = deserializer
@@ -271,6 +553,15 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         self.pipe = pipe
         self.supports_pipeline = Capabilities().has_pipeline()
         self.lock = threading.Lock()
+        self.embedding_config = embedding
+        if self.embedding_config:
+            self.embeddings: Optional[Embeddings] = ensure_embeddings(
+                self.embedding_config.get("embed"),
+                aembed=self.embedding_config.get("aembed"),
+            )
+        else:
+            self.embeddings = None
+        # TODO: Coerce embedding regular functions
 
     @classmethod
     @contextmanager
@@ -280,15 +571,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         *,
         pipeline: bool = False,
         pool_config: Optional[PoolConfig] = None,
+        embedding: Optional[PostgresEmbeddingConfig] = None,
     ) -> Iterator["PostgresStore"]:
         """Create a new PostgresStore instance from a connection string.
 
         Args:
             conn_string (str): The Postgres connection info string.
-            pipeline (bool): whether to use Pipeline (only for single connections)
+            pipeline (bool): whether to use Pipeline
             pool_config (Optional[PoolArgs]): Configuration for the connection pool.
                 If provided, will create a connection pool and use it instead of a single connection.
                 This overrides the `pipeline` argument.
+            embedding (Optional[PostgresEmbeddingConfig]): The embedding config.
+
         Returns:
             PostgresStore: A new PostgresStore instance.
         """
@@ -309,16 +603,16 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                     **cast(dict, pc),
                 ),
             ) as pool:
-                yield cls(conn=pool)
+                yield cls(conn=pool, embedding=embedding)
         else:
             with Connection.connect(
                 conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
             ) as conn:
                 if pipeline:
                     with conn.pipeline() as pipe:
-                        yield cls(conn, pipe=pipe)
+                        yield cls(conn, pipe=pipe, embedding=embedding)
                 else:
-                    yield cls(conn)
+                    yield cls(conn, embedding=embedding)
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -344,14 +638,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:
@@ -414,7 +712,29 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         put_ops: Sequence[tuple[int, PutOp]],
         cur: Cursor[DictRow],
     ) -> None:
-        queries = self._get_batch_PUT_queries(put_ops)
+        queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
+        if embedding_request:
+            if self.embeddings is None:
+                # Should not get here since the embedding config is required
+                # to return an embedding_request above
+                raise ValueError(
+                    "Embedding configuration is required for vector operations "
+                    f"(for semantic search). "
+                    f"Please provide an Embeddings when initializing the {self.__class__.__name__}."
+                )
+            query, txt_params = embedding_request
+            # Update the params to replace the raw text with the vectors
+            vectors = self.embeddings.embed_documents(
+                [param[-1] for param in txt_params]
+            )
+
+            queries.extend(
+                [
+                    (query, (ns, key, value, vector))
+                    for (ns, key, value, _), vector in zip(txt_params, vectors)
+                ]
+            )
+
         for query, params in queries:
             cur.execute(query, params)
 
@@ -424,14 +744,24 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         results: list[Result],
         cur: Cursor[DictRow],
     ) -> None:
-        for (query, params), (idx, _) in zip(
-            self._get_batch_search_queries(search_ops), search_ops
-        ):
+        queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
+
+        if embedding_requests and self.embeddings:
+            embeddings = self.embeddings.embed_documents(
+                [query for _, query in embedding_requests]
+            )
+            for (idx, _), embedding in zip(embedding_requests, embeddings):
+                queries[idx][1][0] = embedding
+
+        for (idx, _), (query, params) in zip(search_ops, queries):
             cur.execute(query, params)
             rows = cast(list[Row], cur.fetchall())
             results[idx] = [
                 _row_to_item(
-                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                    _decode_ns_bytes(row["prefix"]),
+                    row,
+                    loader=self._deserializer,
+                    cls=SearchItem,
                 )
                 for row in rows
             ]
@@ -478,9 +808,25 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             for v, migration in enumerate(
                 self.MIGRATIONS[version + 1 :], start=version + 1
             ):
-                cur.execute(migration)
+                if isinstance(migration, str):
+                    sql = migration
+                else:
+                    if migration.condition and not migration.condition(self):
+                        continue
+
+                    sql = migration.sql
+                    if migration.params:
+                        params = {
+                            k: v(self) if v is not None and callable(v) else v
+                            for k, v in migration.params.items()
+                        }
+                        sql = sql % params
+                cur.execute(sql)
                 cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
 
+        if self.pipe:
+            self.pipe.sync()
+
 
 class Row(TypedDict):
     key: str
@@ -504,17 +850,32 @@ def _row_to_item(
     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
+    cls: Union[type[SearchItem], type[Item]] = Item,
+) -> Union[Item, SearchItem]:
+    """Convert a row from the database into an Item.
+
+    Args:
+        namespace: Item namespace
+        row: Database row
+        loader: Optional value loader for non-dict values
+        cls: Item class to instantiate (Item or SearchItem)
+    """
     val = row["value"]
-    return Item(
-        value=val if isinstance(val, dict) else loader(val),
-        key=row["key"],
-        namespace=namespace,
-        created_at=row["created_at"],
-        updated_at=row["updated_at"],
-    )
+    if not isinstance(val, dict):
+        val = (loader or _json_loads)(val)
+
+    kwargs = {
+        "key": row["key"],
+        "namespace": namespace,
+        "value": val,
+        "created_at": row["created_at"],
+        "updated_at": row["updated_at"],
+    }
+
+    if cls is SearchItem and "score" in row:
+        kwargs["response_metadata"] = {"score": float(row["score"])}
+
+    return cls(**kwargs)
 
 
 def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
@@ -544,3 +905,173 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
     if isinstance(namespace, bytes):
         namespace = namespace.decode()[1:]
     return tuple(namespace.split("."))
+
+
+def _tokenize_path(path: str) -> list[str]:
+    """Tokenize a path into components.
+
+    Handles:
+    - Simple paths: "field1.field2"
+    - Array indexing: "[0]", "[*]", "[-1]"
+    - Wildcards: "*"
+    - Multi-field selection: "{field1,field2}"
+    """
+    if not path:
+        return []
+
+    tokens = []
+    current: list[str] = []
+    i = 0
+    while i < len(path):
+        char = path[i]
+
+        if char == "[":  # Handle array index
+            if current:
+                tokens.append("".join(current))
+                current = []
+            bracket_count = 1
+            index_chars = ["["]
+            i += 1
+            while i < len(path) and bracket_count > 0:
+                if path[i] == "[":
+                    bracket_count += 1
+                elif path[i] == "]":
+                    bracket_count -= 1
+                index_chars.append(path[i])
+                i += 1
+            tokens.append("".join(index_chars))
+            continue
+
+        elif char == "{":  # Handle multi-field selection
+            if current:
+                tokens.append("".join(current))
+                current = []
+            brace_count = 1
+            field_chars = ["{"]
+            i += 1
+            while i < len(path) and brace_count > 0:
+                if path[i] == "{":
+                    brace_count += 1
+                elif path[i] == "}":
+                    brace_count -= 1
+                field_chars.append(path[i])
+                i += 1
+            tokens.append("".join(field_chars))
+            continue
+
+        elif char == ".":  # Handle regular field
+            if current:
+                tokens.append("".join(current))
+                current = []
+        else:
+            current.append(char)
+        i += 1
+
+    if current:
+        tokens.append("".join(current))
+
+    return tokens
+
+
+def _extract_text_by_path(obj: Any, path: str) -> list[str]:
+    """Extract text from an object using a path expression.
+
+    Supports:
+    - Simple paths: "field1.field2"
+    - Array indexing: "[0]", "[*]", "[-1]"
+    - Wildcards: "*"
+    - Multi-field selection: "{field1,field2}"
+    - Nested paths in multi-field: "{field1,nested.field2}"
+    """
+    if not path or path == "__root__":
+        return [json.dumps(obj, sort_keys=True)]
+
+    def _extract_from_obj(obj: Any, tokens: list[str], pos: int) -> list[str]:
+        if pos >= len(tokens):
+            if isinstance(obj, (str, int, float, bool)):
+                return [str(obj)]
+            elif obj is None:
+                return []
+            elif isinstance(obj, (list, dict)):
+                return [json.dumps(obj, sort_keys=True)]
+            return []
+
+        token = tokens[pos]
+        results = []
+
+        if token.startswith("[") and token.endswith("]"):
+            if not isinstance(obj, list):
+                return []
+
+            index = token[1:-1]
+            if index == "*":
+                for item in obj:
+                    results.extend(_extract_from_obj(item, tokens, pos + 1))
+            else:
+                try:
+                    idx = int(index)
+                    if idx < 0:
+                        idx = len(obj) + idx
+                    if 0 <= idx < len(obj):
+                        results.extend(_extract_from_obj(obj[idx], tokens, pos + 1))
+                except (ValueError, IndexError):
+                    return []
+
+        elif token.startswith("{") and token.endswith("}"):
+            if not isinstance(obj, dict):
+                return []
+
+            fields = [f.strip() for f in token[1:-1].split(",")]
+            for field in fields:
+                nested_tokens = _tokenize_path(field)
+                if nested_tokens:
+                    current_obj: Optional[dict] = obj
+                    for nested_token in nested_tokens:
+                        if (
+                            isinstance(current_obj, dict)
+                            and nested_token in current_obj
+                        ):
+                            current_obj = current_obj[nested_token]
+                        else:
+                            current_obj = None
+                            break
+                    if current_obj is not None:
+                        if isinstance(current_obj, (str, int, float, bool)):
+                            results.append(str(current_obj))
+                        elif isinstance(current_obj, (list, dict)):
+                            results.append(json.dumps(current_obj, sort_keys=True))
+
+        # Handle wildcard
+        elif token == "*":
+            if isinstance(obj, dict):
+                for value in obj.values():
+                    results.extend(_extract_from_obj(value, tokens, pos + 1))
+            elif isinstance(obj, list):
+                for item in obj:
+                    results.extend(_extract_from_obj(item, tokens, pos + 1))
+
+        # Handle regular field
+        else:
+            if isinstance(obj, dict) and token in obj:
+                results.extend(_extract_from_obj(obj[token], tokens, pos + 1))
+
+        return results
+
+    tokens = _tokenize_path(path)
+    return _extract_from_obj(obj, tokens, 0)
+
+
+def _get_distance_operator(store: Any) -> tuple[str, str]:
+    """Get the distance operator and score expression based on config."""
+    if not store.embedding_config:
+        return "<=>", "1 - (sv.embedding <=> %s::vector)"
+
+    config = cast(PostgresEmbeddingConfig, store.embedding_config)
+    distance_type = config.get("distance_type", "cosine")
+
+    if distance_type == "l2":
+        return "<->", "1 - (sv.embedding <-> %s::%s)"
+    elif distance_type == "inner_product":
+        return "<#>", "-(sv.embedding <#> %s::%s)"
+    else:  # cosine
+        return "<=>", "1 - (sv.embedding <=> %s::%s)"

libs/checkpoint-postgres/tests/compose-postgres.yml

@@ -1,12 +1,13 @@
 services:
   postgres-test:
-    image: postgres:${POSTGRES_VERSION:-16}
+    image: pgvector/pgvector:pg${POSTGRES_VERSION:-16}
     ports:
       - "5441:5432"
     environment:
       POSTGRES_DB: postgres
       POSTGRES_USER: postgres
       POSTGRES_PASSWORD: postgres
+    command: ["postgres", "-c", "shared_preload_libraries=vector"]
     healthcheck:
       test: pg_isready -U postgres
       start_period: 10s

libs/checkpoint-postgres/tests/conftest.py

@@ -1,9 +1,10 @@
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 import pytest
 from psycopg import AsyncConnection
 from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
+from utils import CharacterEmbeddings  # type: ignore
 
 DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
 
@@ -31,3 +32,12 @@ async def clear_test_db(conn: AsyncConnection[DictRow]) -> None:
         await conn.execute("DELETE FROM store")
     except UndefinedTable:
         pass
+
+
+@pytest.fixture
+def fake_embeddings() -> CharacterEmbeddings:
+    return CharacterEmbeddings(dims=500)
+
+
+INDEX_TYPES = ["hnsw", "ivfflat"]
+VECTOR_TYPES = ["vector", "halfvec"]

libs/checkpoint-postgres/tests/test_async_store.py

@@ -1,10 +1,16 @@
 # type: ignore
 import sys
 import uuid
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 import pytest
-from conftest import DEFAULT_URI  # type: ignore
+from conftest import (
+    DEFAULT_URI,  # type: ignore
+    INDEX_TYPES,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
+from langchain_core.embeddings import Embeddings
 from psycopg import AsyncConnection
 
 from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
@@ -181,272 +187,214 @@ async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
     assert ("test", "namespace2") in results[0]
 
 
-class TestAsyncPostgresStore:
-    @pytest.fixture(autouse=True)
-    async def setup(self) -> None:
-        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
+@pytest.fixture(
+    scope="function",
+    params=[
+        (index_type, vector_type, distance_type)
+        for index_type in INDEX_TYPES
+        for vector_type in VECTOR_TYPES
+        for distance_type in (
+            (["hamming"] if index_type == "ivfflat" else ["hamming", "jaccard"])
+            if vector_type == "bit"
+            else ["l2", "inner_product", "cosine"]
+        )
+    ],
+    ids=lambda p: f"{p[0]}_{p[1]}_{p[2]}",
+)
+async def vector_store(
+    request,
+    fake_embeddings: CharacterEmbeddings,
+) -> AsyncIterator[AsyncPostgresStore]:
+    """Create a store with vector search enabled."""
+    if sys.version_info < (3, 10):
+        pytest.skip("Async Postgres tests require Python 3.10+")
+
+    database = f"test_{uuid.uuid4().hex[:16]}"
+    uri_parts = DEFAULT_URI.split("/")
+    uri_base = "/".join(uri_parts[:-1])
+    query_params = ""
+    if "?" in uri_parts[-1]:
+        db_name, query_params = uri_parts[-1].split("?", 1)
+        query_params = "?" + query_params
+
+    conn_string = f"{uri_base}/{database}{query_params}"
+    admin_conn_string = DEFAULT_URI
+
+    index_type, vector_type, distance_type = request.param
+    embedding_config = {
+        "dims": fake_embeddings.dims,
+        "embed": fake_embeddings,
+        "index_config": {
+            "kind": index_type,
+            "vector_type": vector_type,
+        },
+        "distance_type": distance_type,
+    }
+
+    async with await AsyncConnection.connect(
+        admin_conn_string, autocommit=True
+    ) as conn:
+        await conn.execute(f"CREATE DATABASE {database}")
+    try:
+        async with AsyncPostgresStore.from_conn_string(
+            conn_string,
+            embedding=embedding_config,
+        ) as store:
             await store.setup()
+            yield store
+    finally:
+        async with await AsyncConnection.connect(
+            admin_conn_string, autocommit=True
+        ) as conn:
+            await conn.execute(f"DROP DATABASE {database}")
 
-    async def test_basic_store_ops(self) -> None:
-        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
-            namespace = ("test", "documents")
-            item_id = "doc1"
-            item_value = {"title": "Test Document", "content": "Hello, World!"}
 
-            await store.aput(namespace, item_id, item_value)
-            item = await store.aget(namespace, item_id)
+async def test_vector_store_initialization(
+    vector_store: AsyncPostgresStore, fake_embeddings: CharacterEmbeddings
+) -> None:
+    """Test store initialization with embedding config."""
+    assert vector_store.embedding_config is not None
+    assert vector_store.embedding_config["dims"] == fake_embeddings.dims
+    if isinstance(vector_store.embedding_config["embed"], Embeddings):
+        assert vector_store.embedding_config["embed"] == fake_embeddings
 
-            assert item
-            assert item.namespace == namespace
-            assert item.key == item_id
-            assert item.value == item_value
 
-            updated_value = {
-                "title": "Updated Test Document",
-                "content": "Hello, LangGraph!",
-            }
-            await store.aput(namespace, item_id, updated_value)
-            updated_item = await store.aget(namespace, item_id)
+async def test_vector_insert_with_auto_embedding(
+    vector_store: AsyncPostgresStore,
+) -> None:
+    """Test inserting items that get auto-embedded."""
+    docs = [
+        ("doc1", {"text": "short text"}),
+        ("doc2", {"text": "longer text document"}),
+        ("doc3", {"text": "longest text document here"}),
+        ("doc4", {"description": "text in description field"}),
+        ("doc5", {"content": "text in content field"}),
+        ("doc6", {"body": "text in body field"}),
+    ]
 
-            assert updated_item.value == updated_value
-            assert updated_item.updated_at > item.updated_at
-            different_namespace = ("test", "other_documents")
-            item_in_different_namespace = await store.aget(different_namespace, item_id)
-            assert item_in_different_namespace is None
+    for key, value in docs:
+        await vector_store.aput(("test",), key, value)
 
-            new_item_id = "doc2"
-            new_item_value = {"title": "Another Document", "content": "Greetings!"}
-            await store.aput(namespace, new_item_id, new_item_value)
+    results = await vector_store.asearch(("test",), query="long text")
+    assert len(results) > 0
 
-            search_results = await store.asearch(["test"], limit=10)
-            items = search_results
-            assert len(items) == 2
-            assert any(item.key == item_id for item in items)
-            assert any(item.key == new_item_id for item in items)
+    doc_order = [r.key for r in results]
+    assert "doc2" in doc_order
+    assert "doc3" in doc_order
 
-            namespaces = await store.alist_namespaces(prefix=["test"])
-            assert ("test", "documents") in namespaces
 
-            await store.adelete(namespace, item_id)
-            await store.adelete(namespace, new_item_id)
-            deleted_item = await store.aget(namespace, item_id)
-            assert deleted_item is None
+async def test_vector_update_with_embedding(vector_store: AsyncPostgresStore) -> None:
+    """Test that updating items properly updates their embeddings."""
+    await vector_store.aput(("test",), "doc1", {"text": "zany zebra Xerxes"})
+    await vector_store.aput(("test",), "doc2", {"text": "something about dogs"})
+    await vector_store.aput(("test",), "doc3", {"text": "text about birds"})
 
-            deleted_item = await store.aget(namespace, new_item_id)
-            assert deleted_item is None
+    results_initial = await vector_store.asearch(("test",), query="Zany Xerxes")
+    assert len(results_initial) > 0
+    assert results_initial[0].key == "doc1"
+    initial_score = results_initial[0].response_metadata["score"]
 
-            empty_search_results = await store.asearch(["test"], limit=10)
-            assert len(empty_search_results) == 0
+    await vector_store.aput(("test",), "doc1", {"text": "new text about dogs"})
 
-    async def test_list_namespaces(self) -> None:
-        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
-            test_pref = str(uuid.uuid4())
-            test_namespaces = [
-                (test_pref, "test", "documents", "public", test_pref),
-                (test_pref, "test", "documents", "private", test_pref),
-                (test_pref, "test", "images", "public", test_pref),
-                (test_pref, "test", "images", "private", test_pref),
-                (test_pref, "prod", "documents", "public", test_pref),
-                (
-                    test_pref,
-                    "prod",
-                    "documents",
-                    "some",
-                    "nesting",
-                    "public",
-                    test_pref,
-                ),
-                (test_pref, "prod", "documents", "private", test_pref),
-            ]
+    results_after = await vector_store.asearch(("test",), query="Zany Xerxes")
+    after_score = next(
+        (r.response_metadata["score"] for r in results_after if r.key == "doc1"), 0.0
+    )
+    assert after_score < initial_score
 
-            for namespace in test_namespaces:
-                await store.aput(namespace, "dummy", {"content": "dummy"})
+    results_new = await vector_store.asearch(("test",), query="new text about dogs")
+    for r in results_new:
+        if r.key == "doc1":
+            assert r.response_metadata["score"] > after_score
 
-            prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
-            assert len(prefix_result) == 4
-            assert all([ns[1] == "test" for ns in prefix_result])
+    # Don't index this one
+    await vector_store.aput(
+        ("test",), "doc4", {"text": "new text about dogs"}, index=False
+    )
+    results_new = await vector_store.asearch(
+        ("test",), query="new text about dogs", limit=3
+    )
+    assert not any(r.key == "doc4" for r in results_new)
 
-            specific_prefix_result = await store.alist_namespaces(
-                prefix=[test_pref, "test", "documents"]
-            )
-            assert len(specific_prefix_result) == 2
-            assert all(
-                [ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
-            )
 
-            suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
-            assert len(suffix_result) == 4
-            assert all(ns[-2] == "public" for ns in suffix_result)
+async def test_vector_search_with_filters(vector_store: AsyncPostgresStore) -> None:
+    """Test combining vector search with filters."""
+    docs = [
+        ("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
+        ("doc2", {"text": "red car", "color": "red", "score": 3.0}),
+        ("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
+        ("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
+    ]
 
-            prefix_suffix_result = await store.alist_namespaces(
-                prefix=[test_pref, "test"], suffix=["public", test_pref]
-            )
-            assert len(prefix_suffix_result) == 2
-            assert all(
-                ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
-            )
+    for key, value in docs:
+        await vector_store.aput(("test",), key, value)
 
-            wildcard_prefix_result = await store.alist_namespaces(
-                prefix=[test_pref, "*", "documents"]
-            )
-            assert len(wildcard_prefix_result) == 5
-            assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
+    results = await vector_store.asearch(
+        ("test",), query="apple", filter={"color": "red"}
+    )
+    assert len(results) == 2
+    assert results[0].key == "doc1"
 
-            wildcard_suffix_result = await store.alist_namespaces(
-                suffix=["*", "public", test_pref]
-            )
-            assert len(wildcard_suffix_result) == 4
-            assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
-            wildcard_single = await store.alist_namespaces(
-                suffix=["some", "*", "public", test_pref]
-            )
-            assert len(wildcard_single) == 1
-            assert wildcard_single[0] == (
-                test_pref,
-                "prod",
-                "documents",
-                "some",
-                "nesting",
-                "public",
-                test_pref,
-            )
+    results = await vector_store.asearch(
+        ("test",), query="car", filter={"color": "red"}
+    )
+    assert len(results) == 2
+    assert results[0].key == "doc2"
 
-            max_depth_result = await store.alist_namespaces(max_depth=3)
-            assert all([len(ns) <= 3 for ns in max_depth_result])
-            max_depth_result = await store.alist_namespaces(
-                max_depth=4, prefix=[test_pref, "*", "documents"]
-            )
-            assert (
-                len(set(tuple(res) for res in max_depth_result))
-                == len(max_depth_result)
-                == 5
-            )
+    results = await vector_store.asearch(
+        ("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
+    )
+    assert len(results) == 3
+    assert results[0].key == "doc4"
 
-            limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
-            assert len(limit_result) == 3
+    results = await vector_store.asearch(
+        ("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
+    )
+    assert len(results) == 1
+    assert results[0].key == "doc3"
 
-            offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
-            assert len(offset_result) == len(test_namespaces) - 3
 
-            empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
-            assert len(empty_prefix_result) == len(test_namespaces)
-            assert set(tuple(ns) for ns in empty_prefix_result) == set(
-                tuple(ns) for ns in test_namespaces
-            )
+async def test_vector_search_pagination(vector_store: AsyncPostgresStore) -> None:
+    """Test pagination with vector search."""
+    for i in range(5):
+        await vector_store.aput(
+            ("test",), f"doc{i}", {"text": f"test document number {i}"}
+        )
 
-            for namespace in test_namespaces:
-                await store.adelete(namespace, "dummy")
+    results_page1 = await vector_store.asearch(("test",), query="test", limit=2)
+    results_page2 = await vector_store.asearch(
+        ("test",), query="test", limit=2, offset=2
+    )
 
-    async def test_search(self):
-        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
-            test_namespaces = [
-                ("test_search", "documents", "user1"),
-                ("test_search", "documents", "user2"),
-                ("test_search", "reports", "department1"),
-                ("test_search", "reports", "department2"),
-            ]
-            test_items = [
-                {"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
-                {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
-                {"title": "Report A", "author": "John Doe", "tags": ["final"]},
-                {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
-            ]
-            empty = await store.asearch(
-                (
-                    "scoped",
-                    "assistant_id",
-                    "shared",
-                    "6c5356f6-63ab-4158-868d-cd9fd14c736e",
-                ),
-                limit=10,
-                offset=0,
-            )
-            assert len(empty) == 0
+    assert len(results_page1) == 2
+    assert len(results_page2) == 2
+    assert results_page1[0].key != results_page2[0].key
 
-            for namespace, item in zip(test_namespaces, test_items):
-                await store.aput(namespace, f"item_{namespace[-1]}", item)
+    all_results = await vector_store.asearch(("test",), query="test", limit=10)
+    assert len(all_results) == 5
 
-            docs_result = await store.asearch(["test_search", "documents"])
-            assert len(docs_result) == 2
-            assert all([item.namespace[1] == "documents" for item in docs_result]), [
-                item.namespace for item in docs_result
-            ]
 
-            reports_result = await store.asearch(["test_search", "reports"])
-            assert len(reports_result) == 2
-            assert all(item.namespace[1] == "reports" for item in reports_result)
+async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> None:
+    """Test edge cases in vector search."""
+    await vector_store.aput(("test",), "doc1", {"text": "test document"})
 
-            limited_result = await store.asearch(["test_search"], limit=2)
-            assert len(limited_result) == 2
-            offset_result = await store.asearch(["test_search"])
-            assert len(offset_result) == 4
+    perfect_match = await vector_store.asearch(("test",), query="text test document")
+    perfect_score = perfect_match[0].response_metadata["score"]
 
-            offset_result = await store.asearch(["test_search"], offset=2)
-            assert len(offset_result) == 2
-            assert all(item not in limited_result for item in offset_result)
+    results = await vector_store.asearch(("test",), query="")
+    assert len(results) == 1
+    assert "score" not in results[0].response_metadata
 
-            john_doe_result = await store.asearch(
-                ["test_search"], filter={"author": "John Doe"}
-            )
-            assert len(john_doe_result) == 2
-            assert all(item.value["author"] == "John Doe" for item in john_doe_result)
+    results = await vector_store.asearch(("test",), query=None)
+    assert len(results) == 1
+    assert "score" not in results[0].response_metadata
 
-            draft_result = await store.asearch(
-                ["test_search"], filter={"tags": ["draft"]}
-            )
-            assert len(draft_result) == 2
-            assert all("draft" in item.value["tags"] for item in draft_result)
+    long_query = "foo " * 100
+    results = await vector_store.asearch(("test",), query=long_query)
+    assert len(results) == 1
+    assert results[0].response_metadata["score"] < perfect_score
 
-            page1 = await store.asearch(["test_search"], limit=2, offset=0)
-            page2 = await store.asearch(["test_search"], limit=2, offset=2)
-            all_items = page1 + page2
-            assert len(all_items) == 4
-            assert len(set(item.key for item in all_items)) == 4
-            empty = await store.asearch(
-                (
-                    "scoped",
-                    "assistant_id",
-                    "shared",
-                    "again",
-                    "maybe",
-                    "some-long",
-                    "6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
-                ),
-                limit=10,
-                offset=0,
-            )
-            assert len(empty) == 0
-
-            # Test with a namespace beginning with a number (like a UUID)
-            uuid_namespace = (str(uuid.uuid4()), "documents")
-            uuid_item_id = "uuid_doc"
-            uuid_item_value = {
-                "title": "UUID Document",
-                "content": "This document has a UUID namespace.",
-            }
-
-            # Insert the item with the UUID namespace
-            await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
-
-            # Retrieve the item to verify it was stored correctly
-            retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
-            assert retrieved_item is not None
-            assert retrieved_item.namespace == uuid_namespace
-            assert retrieved_item.key == uuid_item_id
-            assert retrieved_item.value == uuid_item_value
-
-            # Search for the item using the UUID namespace
-            search_result = await store.asearch([uuid_namespace[0]])
-            assert len(search_result) == 1
-            assert search_result[0].key == uuid_item_id
-            assert search_result[0].value == uuid_item_value
-
-            # Clean up: delete the item with the UUID namespace
-            await store.adelete(uuid_namespace, uuid_item_id)
-
-            # Verify the item was deleted
-            deleted_item = await store.aget(uuid_namespace, uuid_item_id)
-            assert deleted_item is None
-
-            for namespace in test_namespaces:
-                await store.adelete(namespace, f"item_{namespace[-1]}")
+    special_query = "test!@#$%^&*()"
+    results = await vector_store.asearch(("test",), query=special_query)
+    assert len(results) == 1
+    assert results[0].response_metadata["score"] < perfect_score

libs/checkpoint-postgres/tests/test_store.py

@@ -1,9 +1,16 @@
 # type: ignore
 
+import json
 from uuid import uuid4
 
 import pytest
-from conftest import DEFAULT_URI  # type: ignore
+from conftest import (
+    DEFAULT_URI,  # type: ignore
+    INDEX_TYPES,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
+from langchain_core.embeddings import Embeddings
 from psycopg import Connection
 
 from langgraph.store.base import (
@@ -15,6 +22,7 @@ from langgraph.store.base import (
     SearchOp,
 )
 from langgraph.store.postgres import PostgresStore
+from langgraph.store.postgres.base import _extract_text_by_path
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -340,3 +348,258 @@ class TestPostgresStore:
             # Cleanup
             for namespace, key, _ in test_data:
                 store.delete(namespace, key)
+
+
+@pytest.fixture(
+    scope="function",
+    params=[
+        (index_type, vector_type, distance_type)
+        for index_type in INDEX_TYPES
+        for vector_type in VECTOR_TYPES
+        for distance_type in (
+            (["hamming"] if index_type == "ivfflat" else ["hamming", "jaccard"])
+            if vector_type == "bit"
+            else ["l2", "inner_product", "cosine"]
+        )
+    ],
+    ids=lambda p: f"{p[0]}_{p[1]}_{p[2]}",
+)
+def vector_store(request, fake_embeddings: Embeddings) -> PostgresStore:
+    """Create a store with vector search enabled."""
+    database = f"test_{uuid4().hex[:16]}"
+    uri_parts = DEFAULT_URI.split("/")
+    uri_base = "/".join(uri_parts[:-1])
+    query_params = ""
+    if "?" in uri_parts[-1]:
+        db_name, query_params = uri_parts[-1].split("?", 1)
+        query_params = "?" + query_params
+
+    conn_string = f"{uri_base}/{database}{query_params}"
+    admin_conn_string = DEFAULT_URI
+
+    index_type, vector_type, distance_type = request.param
+    embedding_config = {
+        "dims": fake_embeddings.dims,
+        "embed": fake_embeddings,
+        "index_config": {
+            "kind": index_type,
+            "vector_type": vector_type,
+        },
+        "distance_type": distance_type,
+    }
+
+    with Connection.connect(admin_conn_string, autocommit=True) as conn:
+        conn.execute(f"CREATE DATABASE {database}")
+    try:
+        with PostgresStore.from_conn_string(
+            conn_string,
+            embedding=embedding_config,
+        ) as store:
+            store.setup()
+            yield store
+    finally:
+        with Connection.connect(admin_conn_string, autocommit=True) as conn:
+            conn.execute(f"DROP DATABASE {database}")
+
+
+def test_vector_store_initialization(
+    vector_store: PostgresStore, fake_embeddings: CharacterEmbeddings
+) -> None:
+    """Test store initialization with embedding config."""
+    # Store should be initialized with embedding config
+    assert vector_store.embedding_config is not None
+    assert vector_store.embedding_config["dims"] == fake_embeddings.dims
+    assert vector_store.embedding_config["embed"] == fake_embeddings
+
+
+def test_vector_insert_with_auto_embedding(vector_store: PostgresStore) -> None:
+    """Test inserting items that get auto-embedded."""
+    docs = [
+        ("doc1", {"text": "short text"}),
+        ("doc2", {"text": "longer text document"}),
+        ("doc3", {"text": "longest text document here"}),
+        ("doc4", {"description": "text in description field"}),
+        ("doc5", {"content": "text in content field"}),
+        ("doc6", {"body": "text in body field"}),
+    ]
+
+    for key, value in docs:
+        vector_store.put(("test",), key, value)
+
+    results = vector_store.search(("test",), query="long text")
+    assert len(results) > 0
+
+    doc_order = [r.key for r in results]
+    assert "doc2" in doc_order
+    assert "doc3" in doc_order
+
+
+def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
+    """Test that updating items properly updates their embeddings."""
+    vector_store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
+    vector_store.put(("test",), "doc2", {"text": "something about dogs"})
+    vector_store.put(("test",), "doc3", {"text": "text about birds"})
+
+    results_initial = vector_store.search(("test",), query="Zany Xerxes")
+    assert len(results_initial) > 0
+    assert results_initial[0].key == "doc1"
+    initial_score = results_initial[0].response_metadata["score"]
+
+    vector_store.put(("test",), "doc1", {"text": "new text about dogs"})
+
+    results_after = vector_store.search(("test",), query="Zany Xerxes")
+    after_score = next(
+        (r.response_metadata["score"] for r in results_after if r.key == "doc1"), 0.0
+    )
+    assert after_score < initial_score
+
+    results_new = vector_store.search(("test",), query="new text about dogs")
+    for r in results_new:
+        if r.key == "doc1":
+            assert r.response_metadata["score"] > after_score
+
+    # Don't index this one
+    vector_store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
+    results_new = vector_store.search(("test",), query="new text about dogs", limit=3)
+    assert not any(r.key == "doc4" for r in results_new)
+
+
+def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
+    """Test combining vector search with filters."""
+    # Insert test documents
+    docs = [
+        ("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
+        ("doc2", {"text": "red car", "color": "red", "score": 3.0}),
+        ("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
+        ("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
+    ]
+
+    for key, value in docs:
+        vector_store.put(("test",), key, value)
+
+    results = vector_store.search(("test",), query="apple", filter={"color": "red"})
+    assert len(results) == 2
+    assert results[0].key == "doc1"
+
+    results = vector_store.search(("test",), query="car", filter={"color": "red"})
+    assert len(results) == 2
+    assert results[0].key == "doc2"
+
+    results = vector_store.search(
+        ("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
+    )
+    assert len(results) == 3
+    assert results[0].key == "doc4"
+
+    # Multiple filters
+    results = vector_store.search(
+        ("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
+    )
+    assert len(results) == 1
+    assert results[0].key == "doc3"
+
+
+def test_vector_search_pagination(vector_store: PostgresStore) -> None:
+    """Test pagination with vector search."""
+    # Insert multiple similar documents
+    for i in range(5):
+        vector_store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
+
+    # Test with different page sizes
+    results_page1 = vector_store.search(("test",), query="test", limit=2)
+    results_page2 = vector_store.search(("test",), query="test", limit=2, offset=2)
+
+    assert len(results_page1) == 2
+    assert len(results_page2) == 2
+    assert results_page1[0].key != results_page2[0].key
+
+    # Get all results
+    all_results = vector_store.search(("test",), query="test", limit=10)
+    assert len(all_results) == 5
+
+
+def test_vector_search_edge_cases(vector_store: PostgresStore) -> None:
+    """Test edge cases in vector search."""
+    vector_store.put(("test",), "doc1", {"text": "test document"})
+
+    results = vector_store.search(("test",), query="")
+    assert len(results) == 1
+
+    results = vector_store.search(("test",), query=None)
+    assert len(results) == 1
+
+    long_query = "test " * 100
+    results = vector_store.search(("test",), query=long_query)
+    assert len(results) == 1
+
+    special_query = "test!@#$%^&*()"
+    results = vector_store.search(("test",), query=special_query)
+    assert len(results) == 1
+
+
+def test_extract_text_by_path():
+    nested_data = {
+        "name": "test",
+        "info": {
+            "age": 25,
+            "tags": ["a", "b", "c"],
+            "metadata": {"created": "2024-01-01", "updated": "2024-01-02"},
+        },
+        "items": [
+            {"id": 1, "value": "first", "tags": ["x", "y"]},
+            {"id": 2, "value": "second", "tags": ["y", "z"]},
+            {"id": 3, "value": "third", "tags": ["z", "w"]},
+        ],
+        "empty": None,
+        "zeros": [0, 0.0, "0"],
+        "empty_list": [],
+        "empty_dict": {},
+    }
+
+    assert _extract_text_by_path(nested_data, "__root__") == [
+        json.dumps(nested_data, sort_keys=True)
+    ]
+
+    assert _extract_text_by_path(nested_data, "name") == ["test"]
+    assert _extract_text_by_path(nested_data, "info.age") == ["25"]
+
+    assert _extract_text_by_path(nested_data, "info.metadata.created") == ["2024-01-01"]
+
+    assert _extract_text_by_path(nested_data, "items[0].value") == ["first"]
+    assert _extract_text_by_path(nested_data, "items[-1].value") == ["third"]
+    assert _extract_text_by_path(nested_data, "items[1].tags[0]") == ["y"]
+
+    values = _extract_text_by_path(nested_data, "items[*].value")
+    assert set(values) == {"first", "second", "third"}
+
+    metadata_dates = _extract_text_by_path(nested_data, "info.metadata.*")
+    assert set(metadata_dates) == {"2024-01-01", "2024-01-02"}
+    name_and_age = _extract_text_by_path(nested_data, "{name,info.age}")
+    assert set(name_and_age) == {"test", "25"}
+
+    item_fields = _extract_text_by_path(nested_data, "items[*].{id,value}")
+    assert set(item_fields) == {"1", "2", "3", "first", "second", "third"}
+
+    all_tags = _extract_text_by_path(nested_data, "items[*].tags[*]")
+    assert set(all_tags) == {"x", "y", "z", "w"}
+
+    assert _extract_text_by_path(None, "any.path") == []
+    assert _extract_text_by_path({}, "any.path") == []
+    assert _extract_text_by_path(nested_data, "") == [
+        json.dumps(nested_data, sort_keys=True)
+    ]
+    assert _extract_text_by_path(nested_data, "nonexistent") == []
+    assert _extract_text_by_path(nested_data, "items[99].value") == []
+    assert _extract_text_by_path(nested_data, "items[*].nonexistent") == []
+
+    assert _extract_text_by_path(nested_data, "empty") == []
+    assert _extract_text_by_path(nested_data, "empty_list") == ["[]"]
+    assert _extract_text_by_path(nested_data, "empty_dict") == ["{}"]
+
+    zeros = _extract_text_by_path(nested_data, "zeros[*]")
+    assert set(zeros) == {"0", "0.0"}
+
+    assert _extract_text_by_path(nested_data, "items[].value") == []
+    assert _extract_text_by_path(nested_data, "items[abc].value") == []
+    assert _extract_text_by_path(nested_data, "{unclosed") == []
+    assert _extract_text_by_path(nested_data, "nested[{invalid}]") == []

libs/checkpoint-postgres/tests/utils.py

@@ -0,0 +1,63 @@
+import math
+import random
+from collections import Counter
+from typing import Any, Optional
+
+from langchain_core.embeddings import Embeddings
+
+
+class CharacterEmbeddings(Embeddings):
+    """Simple character-frequency based embeddings using random projections."""
+
+    def __init__(self, dims: int = 50, seed: int = 42):
+        """Initialize with embedding dimensions and random seed."""
+        self._rng = random.Random(seed)
+        self._char_to_idx: dict[str, int] = {}
+        self._projection: Optional[list[list[float]]] = None
+        self.dims = dims
+
+    def _ensure_projection_matrix(self, texts: list[str]) -> None:
+        """Lazily initialize character mapping and projection matrix."""
+        if self._projection is None:
+            chars = sorted(set("".join(texts)))
+            self._char_to_idx = {c: i for i, c in enumerate(chars)}
+            self._projection = [
+                [self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims)]
+                for _ in range(len(chars))
+            ]
+
+    def _embed_one(self, text: str) -> list[float]:
+        """Embed a single text."""
+        counts = Counter(text)
+        char_vec = [0.0] * len(self._char_to_idx)
+
+        for char, count in counts.items():
+            if char in self._char_to_idx:
+                char_vec[self._char_to_idx[char]] = count
+
+        total = sum(char_vec)
+        if total > 0:
+            char_vec = [v / total for v in char_vec]
+        embedding = [
+            sum(a * b for a, b in zip(char_vec, proj))
+            for proj in zip(*self._projection)
+        ]
+
+        norm = math.sqrt(sum(x * x for x in embedding))
+        if norm > 0:
+            embedding = [x / norm for x in embedding]
+
+        return embedding
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed a list of documents."""
+        self._ensure_projection_matrix(texts)
+        return [self._embed_one(text) for text in texts]
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed a query string."""
+        self._ensure_projection_matrix([text])
+        return self._embed_one(text)
+
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, CharacterEmbeddings) and self.dims == other.dims

libs/checkpoint/langgraph/store/base/__init__.py

@@ -6,7 +6,15 @@ scoped to user IDs, assistant IDs, or other arbitrary namespaces.
 
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Iterable, Literal, NamedTuple, Optional, Union, cast
+from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast
+
+from langchain_core.embeddings import Embeddings
+
+from langgraph.store.base._embed import (
+    AEmbeddingsFunc,
+    EmbeddingsFunc,
+    ensure_embeddings,
+)
 
 
 class Item:
@@ -73,6 +81,52 @@ class Item:
         }
 
 
+class ResponseMetadata(TypedDict, total=False):
+    """Additional metadata about the response/result."""
+
+    score: float
+    """Relevance/similarity score if from a ranked operation."""
+
+
+class SearchItem(Item):
+    """Represents a result item with additional response metadata."""
+
+    __slots__ = "response_metadata"
+
+    def __init__(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+        value: dict[str, Any],
+        created_at: datetime,
+        updated_at: datetime,
+        response_metadata: Optional[ResponseMetadata] = None,
+    ) -> None:
+        """Initialize a result item.
+
+        Args:
+            namespace: Hierarchical path to the item.
+            key: Unique identifier within the namespace.
+            value: The stored value.
+            created_at: When the item was first created.
+            updated_at: When the item was last updated.
+            response_metadata: Optional metadata about the response/result.
+        """
+        super().__init__(
+            value=value,
+            key=key,
+            namespace=namespace,
+            created_at=created_at,
+            updated_at=updated_at,
+        )
+        self.response_metadata = response_metadata or {}
+
+    def dict(self) -> dict:
+        result = super().dict()
+        result["response_metadata"] = self.response_metadata
+        return result
+
+
 class GetOp(NamedTuple):
     """Operation to retrieve an item by namespace and key."""
 
@@ -93,6 +147,8 @@ class SearchOp(NamedTuple):
     """Maximum number of items to return."""
     offset: int = 0
     """Number of items to skip before returning results."""
+    query: Optional[str] = None
+    """The search query for natural language search."""
 
 
 class PutOp(NamedTuple):
@@ -120,6 +176,12 @@ class PutOp(NamedTuple):
       - Values can be of any serializable type
     - If None, it indicates that the item should be deleted
     """
+    index: Optional[bool] = None
+    """Whether to index the item (if supported by the store).
+    
+    Defaults to True if the store supports indexing. This will embed the document
+    so it can be queried using search.
+    """
 
 
 NameSpacePath = tuple[Union[str, Literal["*"]], ...]
@@ -181,6 +243,38 @@ def _validate_namespace(namespace: tuple[str, ...]) -> None:
         )
 
 
+class EmbeddingConfig(TypedDict, total=False):
+    """Configuration for vector embeddings in PostgreSQL store."""
+
+    dims: int
+    """Number of dimensions in the embedding vectors.
+    
+    Common embedding models have the following dimensions:
+        - OpenAI text-embedding-3-large: 256, 1024, or 3072
+        - OpenAI text-embedding-3-small: 512 or 1536
+        - OpenAI text-embedding-ada-002: 1536
+        - Cohere embed-english-v3.0: 1024
+        - Cohere embed-english-light-v3.0: 384
+        - Cohere embed-multilingual-v3.0: 1024
+        - Cohere embed-multilingual-light-v3.0: 384
+    """
+
+    embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc]
+    """Optional function to generate embeddings from text."""
+    aembed: Optional[AEmbeddingsFunc]
+    """Optional asynchronous function to generate embeddings from text.
+    
+    Provide for asynchronous embedding generation if you do not provide
+    an Embeddings object.
+    """
+
+    text_fields: Optional[list[str]]
+    """Fields to extract text from for embedding generation.
+    
+    Defaults to ["__root__"], which embeds the json object as a whole.
+    """
+
+
 class BaseStore(ABC):
     """Abstract base class for persistent key-value stores.
 
@@ -231,6 +325,7 @@ class BaseStore(ABC):
         namespace_prefix: tuple[str, ...],
         /,
         *,
+        query: Optional[str] = None,
         filter: Optional[dict[str, Any]] = None,
         limit: int = 10,
         offset: int = 0,
@@ -239,6 +334,7 @@ class BaseStore(ABC):
 
         Args:
             namespace_prefix: Hierarchical path prefix to search within.
+            query: Optional query for natural language search.
             filter: Key-value pairs to filter results.
             limit: Maximum number of items to return.
             offset: Number of items to skip before returning results.
@@ -246,18 +342,26 @@ class BaseStore(ABC):
         Returns:
             List of items matching the search criteria.
         """
-        return self.batch([SearchOp(namespace_prefix, filter, limit, offset)])[0]
+        return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0]
 
-    def put(self, namespace: tuple[str, ...], key: str, value: dict[str, Any]) -> None:
+    def put(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+        value: dict[str, Any],
+        index: Optional[bool] = None,
+    ) -> None:
         """Store or update an item.
 
         Args:
             namespace: Hierarchical path for the item.
             key: Unique identifier within the namespace.
             value: Dictionary containing the item's data.
+            index: Whether to index the item (if supported by the store).
+                Defaults to True if the store supports indexing.
         """
         _validate_namespace(namespace)
-        self.batch([PutOp(namespace, key, value)])
+        self.batch([PutOp(namespace, key, value, index=index)])
 
     def delete(self, namespace: tuple[str, ...], key: str) -> None:
         """Delete an item.
@@ -336,6 +440,7 @@ class BaseStore(ABC):
         namespace_prefix: tuple[str, ...],
         /,
         *,
+        query: Optional[str] = None,
         filter: Optional[dict[str, Any]] = None,
         limit: int = 10,
         offset: int = 0,
@@ -344,6 +449,7 @@ class BaseStore(ABC):
 
         Args:
             namespace_prefix: Hierarchical path prefix to search within.
+            query: Optional query for natural language search.
             filter: Key-value pairs to filter results.
             limit: Maximum number of items to return.
             offset: Number of items to skip before returning results.
@@ -351,12 +457,18 @@ class BaseStore(ABC):
         Returns:
             List of items matching the search criteria.
         """
-        return (await self.abatch([SearchOp(namespace_prefix, filter, limit, offset)]))[
-            0
-        ]
+        return (
+            await self.abatch(
+                [SearchOp(namespace_prefix, filter, limit, offset, query)]
+            )
+        )[0]
 
     async def aput(
-        self, namespace: tuple[str, ...], key: str, value: dict[str, Any]
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+        value: dict[str, Any],
+        index: Optional[bool] = None,
     ) -> None:
         """Asynchronously store or update an item.
 
@@ -364,9 +476,11 @@ class BaseStore(ABC):
             namespace: Hierarchical path for the item.
             key: Unique identifier within the namespace.
             value: Dictionary containing the item's data.
+            index: Whether to index the item (if supported by the store).
+                Defaults to True if the store supports indexing.
         """
         _validate_namespace(namespace)
-        await self.abatch([PutOp(namespace, key, value)])
+        await self.abatch([PutOp(namespace, key, value, index)])
 
     async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
         """Asynchronously delete an item.
@@ -427,3 +541,19 @@ class BaseStore(ABC):
             offset=offset,
         )
         return (await self.abatch([op]))[0]
+
+
+__all__ = [
+    "BaseStore",
+    "Item",
+    "Op",
+    "PutOp",
+    "GetOp",
+    "SearchOp",
+    "ListNamespacesOp",
+    "MatchCondition",
+    "NameSpacePath",
+    "NamespaceMatchType",
+    "Embeddings",
+    "ensure_embeddings",
+]

libs/checkpoint/langgraph/store/base/_embed.py

@@ -0,0 +1,204 @@
+"""Utilities for working with embedding functions and LangChain's Embeddings interface.
+
+This module provides tools to wrap arbitrary embedding functions (both sync and async)
+into LangChain's Embeddings interface. This enables using custom embedding functions
+with LangChain-compatible tools while maintaining support for both synchronous and
+asynchronous operations.
+"""
+
+import asyncio
+from typing import Any, Awaitable, Callable, List, Optional, Sequence, Union
+
+from langchain_core.embeddings import Embeddings
+
+EmbeddingsFunc = Callable[[Sequence[str]], list[list[float]]]
+"""Type for synchronous embedding functions.
+
+The function should take a sequence of strings and return a list of embeddings,
+where each embedding is a list of floats. The dimensionality of the embeddings
+should be consistent for all inputs.
+"""
+
+AEmbeddingsFunc = Callable[[Sequence[str]], Awaitable[list[list[float]]]]
+"""Type for asynchronous embedding functions.
+
+Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddings.
+"""
+
+
+def ensure_embeddings(
+    embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None],
+    *,
+    aembed: Optional[AEmbeddingsFunc] = None,
+) -> Embeddings:
+    """Ensure that an embedding function conforms to LangChain's Embeddings interface.
+
+    This function wraps arbitrary embedding functions to make them compatible with
+    LangChain's Embeddings interface. It handles both synchronous and asynchronous
+    functions.
+
+    Args:
+        embed: Either an existing Embeddings instance, or a function that converts
+            text to embeddings. If the function is async, it will be used for both
+            sync and async operations.
+        aembed: Optional async function for embeddings. If provided, it will be used
+            for async operations while the sync function is used for sync operations.
+            Must be None if embed is async.
+
+    Returns:
+        An Embeddings instance that wraps the provided function(s).
+
+    Example:
+        >>> def my_embed_fn(texts): return [[0.1, 0.2] for _ in texts]
+        >>> async def my_async_fn(texts): return [[0.1, 0.2] for _ in texts]
+        >>> # Wrap a sync function
+        >>> embeddings = ensure_embeddings(my_embed_fn)
+        >>> # Wrap an async function
+        >>> embeddings = ensure_embeddings(my_async_fn)
+        >>> # Provide both sync and async implementations
+        >>> embeddings = ensure_embeddings(my_embed_fn, aembed=my_async_fn)
+    """
+    if embed is None and aembed is None:
+        raise ValueError("embed or aembed must be provided")
+    if isinstance(embed, Embeddings):
+        return embed
+    return EmbeddingsLambda(embed, afunc=aembed)
+
+
+class EmbeddingsLambda(Embeddings):
+    """Wrapper to convert embedding functions into LangChain's Embeddings interface.
+
+    This class allows arbitrary embedding functions to be used with LangChain-compatible
+    tools. It supports both synchronous and asynchronous operations, and can be
+    initialized with either:
+    1. A synchronous function for both sync/async operations
+    2. An async function for both sync/async operations
+    3. Both sync and async functions for their respective operations
+
+    The embedding functions should convert text into fixed-dimensional vectors that
+    capture the semantic meaning of the text.
+
+    Args:
+        func: Function that converts text to embeddings. Can be sync or async.
+            If async, it will be used for both sync and async operations.
+        afunc: Optional async function for embeddings. If provided, it will be used
+            for async operations while func is used for sync operations.
+            Must be None if func is async.
+
+    Example:
+        >>> def my_embed_fn(texts):
+        ...     # Return 2D embeddings for each text
+        ...     return [[0.1, 0.2] for _ in texts]
+        >>> embeddings = EmbeddingsLambda(my_embed_fn)
+        >>> result = embeddings.embed_query("hello")  # Returns [0.1, 0.2]
+    """
+
+    def __init__(
+        self,
+        func: Union[EmbeddingsFunc, AEmbeddingsFunc, None],
+        afunc: Optional[AEmbeddingsFunc] = None,
+    ) -> None:
+        if _is_async_callable(func):
+            if afunc is not None:
+                raise ValueError(
+                    "afunc must be None if func is async. The async func will be used for both sync and async operations."
+                )
+            self.afunc = func
+        else:
+            self.func = func
+            if afunc is not None:
+                self.afunc = afunc
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed a list of texts into vectors.
+
+        Args:
+            texts: List of texts to convert to embeddings.
+
+        Returns:
+            List of embeddings, one per input text. Each embedding is a list of floats.
+
+        Raises:
+            ValueError: If the instance was initialized with only an async function.
+        """
+        if not hasattr(self, "func"):
+            raise ValueError(
+                "EmbeddingsLambda was initialized with an async function but no sync function. "
+                "Use aembed_documents for async operation or provide a sync function."
+            )
+        return self.func(texts)
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed a single piece of text.
+
+        Args:
+            text: Text to convert to an embedding.
+
+        Returns:
+            Embedding vector as a list of floats.
+
+        Note:
+            This is equivalent to calling embed_documents with a single text
+            and taking the first result.
+        """
+        return self.embed_documents([text])[0]
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Asynchronously embed a list of texts into vectors.
+
+        Args:
+            texts: List of texts to convert to embeddings.
+
+        Returns:
+            List of embeddings, one per input text. Each embedding is a list of floats.
+
+        Note:
+            If no async function was provided, this falls back to the sync implementation.
+        """
+        if not hasattr(self, "afunc"):
+            return await super().aembed_documents(texts)
+        return await self.afunc(texts)
+
+    async def aembed_query(self, text: str) -> List[float]:
+        """Asynchronously embed a single piece of text.
+
+        Args:
+            text: Text to convert to an embedding.
+
+        Returns:
+            Embedding vector as a list of floats.
+
+        Note:
+            This is equivalent to calling aembed_documents with a single text
+            and taking the first result.
+        """
+        if not hasattr(self, "afunc"):
+            return await super().aembed_query(text)
+        return (await self.afunc([text]))[0]
+
+
+def _is_async_callable(
+    func: Any,
+) -> bool:
+    """Check if a function is async.
+
+    This includes both async def functions and classes with async __call__ methods.
+
+    Args:
+        func: Function or callable object to check.
+
+    Returns:
+        True if the function is async, False otherwise.
+    """
+    return (
+        asyncio.iscoroutinefunction(func)
+        or hasattr(func, "__call__")  # noqa: B004
+        and asyncio.iscoroutinefunction(func.__call__)
+    )
+
+
+__all__ = [
+    "ensure_embeddings",
+    "EmbeddingsFunc",
+    "AEmbeddingsFunc",
+]

libs/checkpoint/langgraph/store/base/batch.py

@@ -43,12 +43,13 @@ class AsyncBatchedBaseStore(BaseStore):
         namespace_prefix: tuple[str, ...],
         /,
         *,
+        query: Optional[str] = None,
         filter: Optional[dict[str, Any]] = None,
         limit: int = 10,
         offset: int = 0,
     ) -> list[Item]:
         fut = self._loop.create_future()
-        self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset)
+        self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query)
         return await fut
 
     async def aput(
@@ -56,10 +57,11 @@ class AsyncBatchedBaseStore(BaseStore):
         namespace: tuple[str, ...],
         key: str,
         value: dict[str, Any],
+        index: Optional[bool] = None,
     ) -> None:
         _validate_namespace(namespace)
         fut = self._loop.create_future()
-        self._aqueue[fut] = PutOp(namespace, key, value)
+        self._aqueue[fut] = PutOp(namespace, key, value, index)
         return await fut
 
     async def adelete(