merge

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

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
@@ -385,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 Any, Callable, Optional, Union, cast
 
 import orjson
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
@@ -19,13 +11,22 @@ 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,
+    SearchOp,
+)
 from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.postgres.base import (
     BasePostgresStore,
     PoolConfig,
+    PostgresIndexConfig,
     Row,
     _decode_ns_bytes,
+    _ensure_index_config,
     _group_ops,
     _row_to_item,
     _row_to_search_item,
@@ -35,7 +36,14 @@ logger = logging.getLogger(__name__)
 
 
 class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
-    __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
+    __slots__ = (
+        "_deserializer",
+        "pipe",
+        "lock",
+        "supports_pipeline",
+        "index_config",
+        "embeddings",
+    )
 
     def __init__(
         self,
@@ -45,6 +53,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         deserializer: Optional[
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
+        index: Optional[PostgresIndexConfig] = None,
     ) -> None:
         if isinstance(conn, AsyncConnectionPool) and pipe is not None:
             raise ValueError(
@@ -57,6 +66,12 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         self.lock = asyncio.Lock()
         self.loop = asyncio.get_running_loop()
         self.supports_pipeline = Capabilities().has_pipeline()
+        self.index_config = index
+        if self.index_config:
+            self.embeddings, self.index_config = _ensure_index_config(self.index_config)
+
+        else:
+            self.embeddings = None
 
     async def abatch(self, ops: Iterable[Op]) -> list[Result]:
         grouped_ops, num_ops = _group_ops(ops)
@@ -71,13 +86,114 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
 
         return results
 
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        pool_config: Optional[PoolConfig] = None,
+        index: Optional[PostgresIndexConfig] = None,
+    ) -> AsyncIterator["AsyncPostgresStore"]:
+        """Create a new AsyncPostgresStore instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+            pipeline (bool): Whether to use AsyncPipeline (only for single connections)
+            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.
+            index (Optional[PostgresIndexConfig]): The embedding config.
+
+        Returns:
+            AsyncPostgresStore: A new AsyncPostgresStore instance.
+        """
+        if pool_config is not None:
+            pc = pool_config.copy()
+            async with cast(
+                AsyncConnectionPool[AsyncConnection[DictRow]],
+                AsyncConnectionPool(
+                    conn_string,
+                    min_size=pc.pop("min_size", 1),
+                    max_size=pc.pop("max_size", None),
+                    kwargs={
+                        "autocommit": True,
+                        "prepare_threshold": 0,
+                        "row_factory": dict_row,
+                        **(pc.pop("kwargs", None) or {}),
+                    },
+                    **cast(dict, pc),
+                ),
+            ) as pool:
+                yield cls(conn=pool, index=index)
+        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, index=index)
+                else:
+                    yield cls(conn=conn, index=index)
+
+    async def setup(self) -> None:
+        """Set up the store database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time the store is used.
+        """
+        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
+                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,))
+
     async def _execute_batch(
         self,
         grouped_ops: dict,
         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]),
@@ -132,7 +248,31 @@ 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
+            vectors = await self.embeddings.aembed_documents(
+                [param[-1] for param in txt_params]
+            )
+            queries.append(
+                (
+                    query,
+                    [
+                        p
+                        for (ns, k, pathname, _), vector in zip(txt_params, vectors)
+                        for p in (ns, k, pathname, vector)
+                    ],
+                )
+            )
+
         for query, params in queries:
             await cur.execute(query, params)
 
@@ -142,8 +282,16 @@ 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:
+            vectors = await self.embeddings.aembed_documents(
+                [query for _, query in embedding_requests]
+            )
+            for (idx, _), vector in zip(embedding_requests, vectors):
+                queries[idx][1][0] = vector
+
+        for (idx, _), (query, params) in zip(search_ops, queries):
             await cur.execute(query, params)
             rows = cast(list[Row], await cur.fetchall())
             items = [
@@ -169,129 +317,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,
                 ):
                     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()
-
-    @classmethod
-    @asynccontextmanager
-    async def from_conn_string(
-        cls,
-        conn_string: str,
-        *,
-        pipeline: bool = False,
-        pool_config: Optional[PoolConfig] = None,
-    ) -> AsyncIterator["AsyncPostgresStore"]:
-        """Create a new AsyncPostgresStore instance from a connection string.
-
-        Args:
-            conn_string (str): The Postgres connection info string.
-            pipeline (bool): Whether to use AsyncPipeline (only for single connections)
-            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.
-
-        Returns:
-            AsyncPostgresStore: A new AsyncPostgresStore instance.
-        """
-        if pool_config is not None:
-            pc = pool_config.copy()
-            async with cast(
-                AsyncConnectionPool[AsyncConnection[DictRow]],
-                AsyncConnectionPool(
-                    conn_string,
-                    min_size=pc.pop("min_size", 1),
-                    max_size=pc.pop("max_size", None),
-                    kwargs={
-                        "autocommit": True,
-                        "prepare_threshold": 0,
-                        "row_factory": dict_row,
-                        **(pc.pop("kwargs", None) or {}),
-                    },
-                    **cast(dict, pc),
-                ),
-            ) as pool:
-                yield cls(conn=pool)
-        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)
-                else:
-                    yield cls(conn=conn)
-
-    async def setup(self) -> None:
-        """Set up the store database asynchronously.
-
-        This method creates the necessary tables in the Postgres database if they don't
-        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:
-                    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,)
-                    )
-            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,
@@ -31,6 +32,7 @@ from langgraph.checkpoint.postgres import _internal as _pg_internal
 from langgraph.store.base import (
     BaseStore,
     GetOp,
+    IndexConfig,
     Item,
     ListNamespacesOp,
     Op,
@@ -38,12 +40,31 @@ from langgraph.store.base import (
     Result,
     SearchItem,
     SearchOp,
+    ensure_embeddings,
+    get_text_at_path,
+    tokenize_path,
 )
 
+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.index_config)
+
+
+MIGRATIONS: Sequence[Union[str, Migration]] = [
     """
 CREATE TABLE IF NOT EXISTS store (
     -- 'prefix' represents the doc's 'namespace'
@@ -59,8 +80,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.index_config["dims"],
+            "vector_type": lambda store: (
+                cast(PostgresIndexConfig, store.index_config)
+                .get("db_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])
 
 
@@ -89,10 +158,76 @@ class PoolConfig(TypedDict, total=False):
     """
 
 
+class DBIndexConfig(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(DBIndexConfig, 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(DBIndexConfig, 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 PostgresIndexConfig(IndexConfig, total=False):
+    """Configuration for vector embeddings in PostgreSQL store with pgvector-specific options.
+
+    Extends EmbeddingConfig with additional configuration for pgvector index and vector types.
+    """
+
+    db_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]]]
+    index_config: Optional[PostgresIndexConfig]
+
+    @staticmethod
+    def _get_default_index_config() -> IndexConfig:
+        return HNSWConfig(
+            kind="hnsw",
+            vector_type="vector",
+        )
 
     def _get_batch_GET_ops_queries(
         self,
@@ -114,10 +249,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:
@@ -144,60 +282,160 @@ 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.index_config:
+                paths = self.index_config["__tokenized_fields"]
+                for op in inserts:
+                    if op.index is False:
+                        continue
+                    value = op.value
+                    ns = _namespace_to_text(op.namespace)
+                    k = op.key
+
+                    for path, tokenized_path in paths:
+                        texts = get_text_at_path(value, tokenized_path)
+                        for i, text in enumerate(texts):
+                            pathname = f"{path}.{i}" if len(texts) > 1 else path
+                            vector_values.append(
+                                "(%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
+                            )
+                            embedding_request_params.append((ns, k, pathname, 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.index_config:
+                needs_vector_search = True
+                embedding_requests.append((idx, op.query))
+
+                score_expr = _get_distance_operator(self)
+                vector_type = (
+                    cast(PostgresIndexConfig, self.index_config)
+                    .get("db_index_config", self._get_default_index_config())
+                    .get("vector_type", "vector")
+                )
+
+                if (
+                    vector_type == "bit"
+                    and self.index_config.get("distance_type") == "hamming"
+                ):
+                    score_expr = score_expr % ("%s", self.index_config["dims"])
+                else:
+                    score_expr = score_expr % ("%s", vector_type)
+
+                vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
+                expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
+
+                # Direct query with DISTINCT ON to get best score per document
+                base_query = f"""
+                    with scored as (
+                        SELECT DISTINCT ON (s.prefix, s.key)
+                            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
+                        ORDER BY s.prefix, s.key, score DESC
+                        LIMIT %s
+                    )
+
+                    SELECT * FROM scored
+                """
+                params = [
+                    None,  # Vector placeholder
+                    f"{_namespace_to_text(op.namespace_prefix)}%",
+                    expanded_limit,
+                ]
 
             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:
+                    if needs_vector_search:
+                        base_query += " WHERE " + " AND ".join(filter_conditions)
+                    else:
+                        base_query += " AND " + " AND ".join(filter_conditions)
+
+            if needs_vector_search:
+                base_query += " ORDER BY score DESC"
+            else:
+                base_query += " ORDER BY updated_at DESC"
+
+            base_query += " 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,
@@ -249,13 +487,37 @@ 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")
+    __slots__ = (
+        "_deserializer",
+        "pipe",
+        "lock",
+        "supports_pipeline",
+        "index_config",
+        "embeddings",
+    )
 
     def __init__(
         self,
@@ -265,6 +527,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         deserializer: Optional[
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
+        index: Optional[PostgresIndexConfig] = None,
     ) -> None:
         super().__init__()
         self._deserializer = deserializer
@@ -272,6 +535,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         self.pipe = pipe
         self.supports_pipeline = Capabilities().has_pipeline()
         self.lock = threading.Lock()
+        self.index_config = index
+        if self.index_config:
+            self.embeddings, self.index_config = _ensure_index_config(self.index_config)
+        else:
+            self.embeddings = None
 
     @classmethod
     @contextmanager
@@ -281,15 +549,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         *,
         pipeline: bool = False,
         pool_config: Optional[PoolConfig] = None,
+        index: Optional[PostgresIndexConfig] = 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[PostgresIndexConfig]): The embedding config.
+
         Returns:
             PostgresStore: A new PostgresStore instance.
         """
@@ -310,16 +581,16 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                     **cast(dict, pc),
                 ),
             ) as pool:
-                yield cls(conn=pool)
+                yield cls(conn=pool, index=index)
         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, index=index)
                 else:
-                    yield cls(conn)
+                    yield cls(conn, index=index)
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -419,7 +690,33 @@ 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.append(
+                (
+                    query,
+                    [
+                        p
+                        for (ns, k, pathname, _), vector in zip(txt_params, vectors)
+                        for p in (ns, k, pathname, vector)
+                    ],
+                )
+            )
+
         for query, params in queries:
             cur.execute(query, params)
 
@@ -429,9 +726,16 @@ 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] = [
@@ -483,7 +787,20 @@ 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,))
 
 
@@ -495,6 +812,56 @@ class Row(TypedDict):
     updated_at: datetime
 
 
+# Private utilities
+
+
+def _get_vector_type_ops(store: BasePostgresStore) -> str:
+    """Get the vector type operator class based on config."""
+    if not store.index_config:
+        return "vector_cosine_ops"
+
+    config = cast(PostgresIndexConfig, store.index_config)
+    index_config = config.get(
+        "db_index_config", BasePostgresStore._get_default_index_config()
+    )
+    vector_type = cast(str, index_config.get("vector_type", "vector"))
+    if vector_type not in ("vector", "halfvec"):
+        raise ValueError(
+            f"Vector type must be 'vector' or 'halfvec', got {vector_type}"
+        )
+
+    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.index_config:
+        return "hnsw", {}
+
+    config = cast(PostgresIndexConfig, store.index_config)
+    default_config = BasePostgresStore._get_default_index_config()
+    index_config = config.get("db_index_config", default_config).copy()
+    kind = index_config.pop("kind", "hnsw")
+    index_config.pop("vector_type", None)
+    return kind, index_config
+
+
 def _namespace_to_text(
     namespace: tuple[str, ...], handle_wildcards: bool = False
 ) -> str:
@@ -510,15 +877,51 @@ def _row_to_item(
     *,
     loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
 ) -> Item:
+    """Convert a row from the database into an Item.
+
+    Args:
+        namespace: Item namespace
+        row: Database row
+        loader: Optional value loader for non-dict values
+    """
+    val = row["value"]
+    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"],
+    }
+
+    return Item(**kwargs)
+
+
+def _row_to_search_item(
+    namespace: tuple[str, ...],
+    row: Row,
+    *,
+    loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
+) -> SearchItem:
     """Convert a row from the database into an Item."""
     loader = loader or _json_loads
     val = row["value"]
-    return Item(
+    score = row.get("score")
+    if score is not None:
+        try:
+            score = float(score)  # type: ignore[arg-type]
+        except ValueError:
+            logger.warning("Invalid score: %s", score)
+            score = None
+    return SearchItem(
         value=val if isinstance(val, dict) else loader(val),
         key=row["key"],
         namespace=namespace,
         created_at=row["created_at"],
         updated_at=row["updated_at"],
+        score=score,
     )
 
 
@@ -575,3 +978,50 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
     if isinstance(namespace, bytes):
         namespace = namespace.decode()[1:]
     return tuple(namespace.split("."))
+
+
+def _get_distance_operator(store: Any) -> str:
+    """Get the distance operator and score expression based on config."""
+    if not store.index_config:
+        raise ValueError(
+            "Embedding configuration is required for vector operations "
+            f"(for semantic search). "
+            f"Please provide an Embeddings when initializing the {store.__class__.__name__}."
+        )
+
+    config = cast(PostgresIndexConfig, store.index_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)"
+
+
+def _ensure_index_config(
+    index_config: PostgresIndexConfig,
+) -> tuple[Optional["Embeddings"], PostgresIndexConfig]:
+    index_config = index_config.copy()
+    tokenized: list[tuple[str, Union[Literal["$"], list[str]]]] = []
+    tot = 0
+    text_fields = index_config.get("text_fields") or ["$"]
+    if isinstance(text_fields, str):
+        text_fields = [text_fields]
+    if not isinstance(text_fields, list):
+        raise ValueError(f"Text fields must be a list or a string. Got {text_fields}")
+    for p in text_fields:
+        if p == "$":
+            tokenized.append((p, "$"))
+            tot += 1
+        else:
+            toks = tokenize_path(p)
+            tokenized.append((p, toks))
+            tot += len(toks)
+    index_config["__tokenized_fields"] = tokenized
+    index_config["__estimated_num_vectors"] = tot
+    embeddings = ensure_embeddings(
+        index_config.get("embed"),
+    )
+    return embeddings, index_config

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,10 +1,12 @@
-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 tests.embed_test_utils import CharacterEmbeddings
+
 DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
 
 
@@ -31,3 +33,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/embed_test_utils.py

@@ -0,0 +1,55 @@
+"""Embedding utilities for testing."""
+
+import math
+import random
+from collections import Counter, defaultdict
+from typing import Any
+
+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.dims = dims
+        # Create projection vector for each character lazily
+        self._char_projections: defaultdict[str, list[float]] = defaultdict(
+            lambda: [
+                self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims)
+            ]
+        )
+
+    def _embed_one(self, text: str) -> list[float]:
+        """Embed a single text."""
+        counts = Counter(text)
+        total = sum(counts.values())
+
+        if total == 0:
+            return [0.0] * self.dims
+
+        embedding = [0.0] * self.dims
+        for char, count in counts.items():
+            weight = count / total
+            char_proj = self._char_projections[char]
+            for i, proj in enumerate(char_proj):
+                embedding[i] += weight * proj
+
+        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."""
+        return [self._embed_one(text) for text in texts]
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed a query string."""
+        return self._embed_one(text)
+
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, CharacterEmbeddings) and self.dims == other.dims

libs/checkpoint-postgres/tests/test_async.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 import pytest
-from conftest import DEFAULT_URI  # type: ignore
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
@@ -11,6 +10,7 @@ from langgraph.checkpoint.base import (
     empty_checkpoint,
 )
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from tests.conftest import DEFAULT_URI
 
 
 class TestAsyncPostgresSaver:

libs/checkpoint-postgres/tests/test_async_store.py

@@ -1,14 +1,22 @@
 # type: ignore
 import sys
 import uuid
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any, Optional
 
 import pytest
-from conftest import DEFAULT_URI  # type: ignore
+from langchain_core.embeddings import Embeddings
 from psycopg import AsyncConnection
 
 from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
 from langgraph.store.postgres import AsyncPostgresStore
+from tests.conftest import (
+    DEFAULT_URI,
+    INDEX_TYPES,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -181,272 +189,292 @@ 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:
+@asynccontextmanager
+async def _create_vector_store(
+    index_type: str,
+    vector_type: str,
+    distance_type: str,
+    fake_embeddings: CharacterEmbeddings,
+    text_fields: Optional[list[str]] = None,
+) -> 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_config = {
+        "dims": fake_embeddings.dims,
+        "embed": fake_embeddings,
+        "db_index_config": {
+            "kind": index_type,
+            "vector_type": vector_type,
+        },
+        "distance_type": distance_type,
+        "text_fields": text_fields,
+    }
+
+    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,
+            index=index_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)
+@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."""
+    index_type, vector_type, distance_type = request.param
+    async with _create_vector_store(
+        index_type, vector_type, distance_type, fake_embeddings
+    ) as store:
+        yield store
 
-            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_store_initialization(
+    vector_store: AsyncPostgresStore, fake_embeddings: CharacterEmbeddings
+) -> None:
+    """Test store initialization with embedding config."""
+    assert vector_store.index_config is not None
+    assert vector_store.index_config["dims"] == fake_embeddings.dims
+    if isinstance(vector_store.index_config["embed"], Embeddings):
+        assert vector_store.index_config["embed"] == fake_embeddings
 
-            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
 
-            new_item_id = "doc2"
-            new_item_value = {"title": "Another Document", "content": "Greetings!"}
-            await store.aput(namespace, new_item_id, new_item_value)
+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"}),
+    ]
 
-            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)
+    for key, value in docs:
+        await vector_store.aput(("test",), key, value)
 
-            namespaces = await store.alist_namespaces(prefix=["test"])
-            assert ("test", "documents") in namespaces
+    results = await vector_store.asearch(("test",), query="long text")
+    assert len(results) > 0
 
-            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
+    doc_order = [r.key for r in results]
+    assert "doc2" in doc_order
+    assert "doc3" in doc_order
 
-            deleted_item = await store.aget(namespace, new_item_id)
-            assert deleted_item is None
 
-            empty_search_results = await store.asearch(["test"], limit=10)
-            assert len(empty_search_results) == 0
+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"})
 
-    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_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].score
 
-            for namespace in test_namespaces:
-                await store.aput(namespace, "dummy", {"content": "dummy"})
+    await vector_store.aput(("test",), "doc1", {"text": "new text about dogs"})
 
-            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])
+    results_after = await vector_store.asearch(("test",), query="Zany Xerxes")
+    after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
+    assert after_score < initial_score
 
-            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]
-            )
+    results_new = await vector_store.asearch(("test",), query="new text about dogs")
+    for r in results_new:
+        if r.key == "doc1":
+            assert r.score > after_score
 
-            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)
+    # 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)
 
-            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
-            )
 
-            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)
+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}),
+    ]
 
-            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,
-            )
+    for key, value in docs:
+        await vector_store.aput(("test",), key, value)
 
-            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="apple", filter={"color": "red"}
+    )
+    assert len(results) == 2
+    assert results[0].key == "doc1"
 
-            limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
-            assert len(limit_result) == 3
+    results = await vector_store.asearch(
+        ("test",), query="car", filter={"color": "red"}
+    )
+    assert len(results) == 2
+    assert results[0].key == "doc2"
 
-            offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
-            assert len(offset_result) == len(test_namespaces) - 3
+    results = await vector_store.asearch(
+        ("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
+    )
+    assert len(results) == 3
+    assert results[0].key == "doc4"
 
-            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
-            )
+    results = await vector_store.asearch(
+        ("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
+    )
+    assert len(results) == 1
+    assert results[0].key == "doc3"
 
-            for namespace in test_namespaces:
-                await store.adelete(namespace, "dummy")
 
-    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
+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, item in zip(test_namespaces, test_items):
-                await store.aput(namespace, f"item_{namespace[-1]}", item)
+    results_page1 = await vector_store.asearch(("test",), query="test", limit=2)
+    results_page2 = await vector_store.asearch(
+        ("test",), query="test", limit=2, offset=2
+    )
 
-            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
-            ]
+    assert len(results_page1) == 2
+    assert len(results_page2) == 2
+    assert results_page1[0].key != results_page2[0].key
 
-            reports_result = await store.asearch(["test_search", "reports"])
-            assert len(reports_result) == 2
-            assert all(item.namespace[1] == "reports" for item in reports_result)
+    all_results = await vector_store.asearch(("test",), query="test", limit=10)
+    assert len(all_results) == 5
 
-            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
 
-            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)
+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"})
 
-            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)
+    perfect_match = await vector_store.asearch(("test",), query="text test document")
+    perfect_score = perfect_match[0].score
 
-            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)
+    results = await vector_store.asearch(("test",), query="")
+    assert len(results) == 1
+    assert results[0].score is None
 
-            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
+    results = await vector_store.asearch(("test",), query=None)
+    assert len(results) == 1
+    assert results[0].score is None
 
-            # 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.",
-            }
+    long_query = "foo " * 100
+    results = await vector_store.asearch(("test",), query=long_query)
+    assert len(results) == 1
+    assert results[0].score < perfect_score
 
-            # Insert the item with the UUID namespace
-            await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
+    special_query = "test!@#$%^&*()"
+    results = await vector_store.asearch(("test",), query=special_query)
+    assert len(results) == 1
+    assert results[0].score < perfect_score
 
-            # 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
+@pytest.mark.parametrize(
+    "index_type,vector_type,distance_type",
+    [
+        ("ivfflat", "vector", "cosine"),
+        ("hnsw", "vector", "cosine"),
+        ("hnsw", "halfvec", "cosine"),
+        ("hnsw", "halfvec", "inner_product"),
+    ],
+)
+async def test_embed_with_path(
+    request: Any,
+    fake_embeddings: CharacterEmbeddings,
+    index_type: str,
+    vector_type: str,
+    distance_type: str,
+) -> None:
+    """Test vector search with specific text fields in Postgres store."""
+    async with _create_vector_store(
+        index_type,
+        vector_type,
+        distance_type,
+        fake_embeddings,
+        text_fields=["key0", "key1", "key3"],
+    ) as store:
+        # This will have 2 vectors representing it
+        doc1 = {
+            # Omit key0 - check it doesn't raise an error
+            "key1": "xxx",
+            "key2": "yyy",
+            "key3": "zzz",
+        }
+        # This will have 3 vectors representing it
+        doc2 = {
+            "key0": "uuu",
+            "key1": "vvv",
+            "key2": "www",
+            "key3": "xxx",
+        }
+        await store.aput(("test",), "doc1", doc1)
+        await store.aput(("test",), "doc2", doc2)
 
-            # Clean up: delete the item with the UUID namespace
-            await store.adelete(uuid_namespace, uuid_item_id)
+        # doc2.key3 and doc1.key1 both would have the highest score
+        results = await store.asearch(("test",), query="xxx")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        ascore = results[0].score
+        bscore = results[1].score
+        assert ascore == pytest.approx(bscore, abs=1e-3)
 
-            # Verify the item was deleted
-            deleted_item = await store.aget(uuid_namespace, uuid_item_id)
-            assert deleted_item is None
+        results = await store.asearch(("test",), query="uuu")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        assert results[0].key == "doc2"
+        assert results[0].score > results[1].score
+        assert ascore == pytest.approx(results[0].score, abs=1e-3)
 
-            for namespace in test_namespaces:
-                await store.adelete(namespace, f"item_{namespace[-1]}")
+        # Un-indexed - will have low results for both. Not zero (because we're projecting)
+        # but less than the above.
+        results = await store.asearch(("test",), query="www")
+        assert len(results) == 2
+        assert results[0].score < ascore
+        assert results[1].score < ascore

libs/checkpoint-postgres/tests/test_store.py

@@ -1,9 +1,11 @@
 # type: ignore
 
+from contextlib import contextmanager
+from typing import Any, Optional
 from uuid import uuid4
 
 import pytest
-from conftest import DEFAULT_URI  # type: ignore
+from langchain_core.embeddings import Embeddings
 from psycopg import Connection
 
 from langgraph.store.base import (
@@ -15,6 +17,12 @@ from langgraph.store.base import (
     SearchOp,
 )
 from langgraph.store.postgres import PostgresStore
+from tests.conftest import (
+    DEFAULT_URI,
+    INDEX_TYPES,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -340,3 +348,281 @@ class TestPostgresStore:
             # Cleanup
             for namespace, key, _ in test_data:
                 store.delete(namespace, key)
+
+
+@contextmanager
+def _create_vector_store(
+    index_type: str,
+    vector_type: str,
+    distance_type: str,
+    fake_embeddings: Embeddings,
+    text_fields: Optional[list[str]] = None,
+) -> 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_config = {
+        "dims": fake_embeddings.dims,
+        "embed": fake_embeddings,
+        "db_index_config": {
+            "kind": index_type,
+            "vector_type": vector_type,
+        },
+        "distance_type": distance_type,
+        "text_fields": text_fields,
+    }
+
+    with Connection.connect(admin_conn_string, autocommit=True) as conn:
+        conn.execute(f"CREATE DATABASE {database}")
+    try:
+        with PostgresStore.from_conn_string(
+            conn_string,
+            index=index_config,
+        ) as store:
+            store.setup()
+            yield store
+    finally:
+        with Connection.connect(admin_conn_string, autocommit=True) as conn:
+            conn.execute(f"DROP DATABASE {database}")
+
+
+@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."""
+    index_type, vector_type, distance_type = request.param
+    with _create_vector_store(
+        index_type, vector_type, distance_type, fake_embeddings
+    ) as store:
+        yield store
+
+
+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.index_config is not None
+    assert vector_store.index_config["dims"] == fake_embeddings.dims
+    assert vector_store.index_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].score
+
+    vector_store.put(("test",), "doc1", {"text": "new text about dogs"})
+
+    results_after = vector_store.search(("test",), query="Zany Xerxes")
+    after_score = next((r.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.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
+
+
+@pytest.mark.parametrize(
+    "index_type,vector_type,distance_type",
+    [
+        ("ivfflat", "vector", "cosine"),
+        ("hnsw", "vector", "cosine"),
+        ("hnsw", "halfvec", "cosine"),
+        ("hnsw", "halfvec", "inner_product"),
+    ],
+)
+def test_embed_with_path_sync(
+    request: Any,
+    fake_embeddings: CharacterEmbeddings,
+    index_type: str,
+    vector_type: str,
+    distance_type: str,
+) -> None:
+    """Test vector search with specific text fields in Postgres store."""
+    with _create_vector_store(
+        index_type,
+        vector_type,
+        distance_type,
+        fake_embeddings,
+        text_fields=["key0", "key1", "key3"],
+    ) as store:
+        # This will have 2 vectors representing it
+        doc1 = {
+            # Omit key0 - check it doesn't raise an error
+            "key1": "xxx",
+            "key2": "yyy",
+            "key3": "zzz",
+        }
+        # This will have 3 vectors representing it
+        doc2 = {
+            "key0": "uuu",
+            "key1": "vvv",
+            "key2": "www",
+            "key3": "xxx",
+        }
+        store.put(("test",), "doc1", doc1)
+        store.put(("test",), "doc2", doc2)
+
+        # doc2.key3 and doc1.key1 both would have the highest score
+        results = store.search(("test",), query="xxx")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        ascore = results[0].score
+        bscore = results[1].score
+        assert ascore == pytest.approx(bscore, abs=1e-3)
+
+        # ~Only match doc2
+        results = store.search(("test",), query="uuu")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        assert results[0].key == "doc2"
+        assert results[0].score > results[1].score
+        assert ascore == pytest.approx(results[0].score, abs=1e-3)
+
+        # ~Only match doc1
+        results = store.search(("test",), query="zzz")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        assert results[0].key == "doc1"
+        assert results[0].score > results[1].score
+        assert ascore == pytest.approx(results[0].score, abs=1e-3)
+
+        # Un-indexed - will have low results for both. Not zero (because we're projecting)
+        # but less than the above.
+        results = store.search(("test",), query="www")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        assert results[0].score < ascore
+        assert results[1].score < ascore

libs/checkpoint-postgres/tests/test_sync.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 import pytest
-from conftest import DEFAULT_URI  # type: ignore
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
@@ -11,6 +10,7 @@ from langgraph.checkpoint.base import (
     empty_checkpoint,
 )
 from langgraph.checkpoint.postgres import PostgresSaver
+from tests.conftest import DEFAULT_URI
 
 
 class TestPostgresSaver: