[SDK] Add SearchItem (#2567)

- Initializing the store with an 'embedding config' -> this contains the
'dims' (used to create the table) and the encoder object (rn langchain
embeddings object, though that is ......)
- Call setup() -> creates the vector table.

Each document has 1 or more vectors associated with it for each json
path in the embedding config.

Would welcome critique and requests! 

Leaving the params as the defaults for pgvector but open to feedback if
you think it's important to be able to more transparently configure that
in setup()

```python
from typing import TypedDict, List, Dict, Any, Optional

from langchain_openai import OpenAIEmbeddings
from langgraph.graph import StateGraph
from langgraph.store.postgres import PostgresStore

emb_config = {
    "dims": 1536,  # OpenAI embedding dimensions
    "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
    "distance_type": "cosine",
}
with PostgresStore.from_conn_string(
    "postgres://postgres:postgres@localhost:5441",
    embedding=emb_config,
) as store:
    store.setup()


# Define the state type for our graph
class State(TypedDict):
    query: str
    results: Optional[List[Dict[str, Any]]]


def put_stuff(state: State) -> State:
    docs = [
        ("doc1", {"text": "red apple in kitchen"}),
        ("doc2", {"text": "blue car in garage"}),
        ("doc3", {"text": "green apple on table"}),
    ]
    for key, value in docs:
        store.put(("docs",), key, value)


def search_stuff(state: State) -> State:
    """Search for documents using vector similarity."""
    results = store.search(("docs",), query=state["query"])

    return {"results": results}


builder = StateGraph(State)
builder.add_node(put_stuff)
builder.add_node(search_stuff)
builder.add_edge("__start__", "put_stuff")
builder.add_edge("put_stuff", "search_stuff")
# Compile
with PostgresStore.from_conn_string(
    "postgres://postgres:postgres@localhost:5441",
    embedding=emb_config,
) as store:
    chain = builder.compile(store=store)

    result = chain.invoke({"query": "sour apple"})

# Print results
for doc in result["results"]:
    print(doc.key)
    print(doc.value)
    print(doc.response_metadata)

```

docs/docs/tutorials/introduction.ipynb

@@ -19,7 +19,7 @@
     "\n",
     "## Setup\n",
     "\n",
-    "First, install the required packages:"
+    "First, install the required packages and configure your environment:"
    ]
   },
   {
@@ -33,14 +33,6 @@
     "%pip install -U langgraph langsmith langchain_anthropic"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "a6d1e870-1bc0-4d44-86c0-96681ccf6113",
-   "metadata": {},
-   "source": [
-    "In this tutorial, we'll be "
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 2,

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,23 @@ 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 (
+    _PLACEHOLDER,
     BasePostgresStore,
     PoolConfig,
+    PostgresIndexConfig,
     Row,
     _decode_ns_bytes,
+    _ensure_index_config,
     _group_ops,
     _row_to_item,
     _row_to_search_item,
@@ -35,7 +37,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 +54,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 +67,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 +87,117 @@ 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 def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
+            try:
+                await cur.execute(f"SELECT v FROM {table} 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(
+                    f"""
+                    CREATE TABLE IF NOT EXISTS {table} (
+                        v INTEGER PRIMARY KEY
+                    )
+                """
+                )
+            return version
+
+        async with self._cursor() as cur:
+            version = await _get_version(cur, table="store_migrations")
+            for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
+                await cur.execute(sql)
+                await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+
+            if self.index_config:
+                version = await _get_version(cur, table="vector_migrations")
+                for v, migration in enumerate(
+                    self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
+                ):
+                    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 vector_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 +252,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 +286,19 @@ 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):
+                _paramslist = queries[idx][1]
+                for i in range(len(_paramslist)):
+                    if _paramslist[i] is _PLACEHOLDER:
+                        _paramslist[i] = 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 +324,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,25 @@ 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
+    params: Optional[dict[str, Any]] = None
+
+
+MIGRATIONS: Sequence[str] = [
     """
 CREATE TABLE IF NOT EXISTS store (
     -- 'prefix' represents the doc's 'namespace'
@@ -61,6 +76,39 @@ CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pa
 """,
 ]
 
+VECTOR_MIGRATIONS: Sequence[Migration] = [
+    Migration(
+        """
+CREATE EXTENSION IF NOT EXISTS vector;
+""",
+    ),
+    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
+);
+""",
+        params={
+            "dims": lambda store: store.index_config["dims"],
+            "vector_type": lambda store: (
+                cast(PostgresIndexConfig, store.index_config)
+                .get("ann_index_config", {})
+                .get("vector_type", "vector")
+            ),
+        },
+    ),
+    # TODO: Add an HNSW or IVFFlat index depending on config
+    # First must improve the search query when filtering by
+    # namespace
+]
+
 C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
 
 
@@ -89,10 +137,39 @@ class PoolConfig(TypedDict, total=False):
     """
 
 
+class ANNIndexConfig(TypedDict, total=False):
+    """Configuration for vector index in PostgreSQL store."""
+
+    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 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.
+    """
+
+    ann_index_config: ANNIndexConfig
+    """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
+    VECTOR_MIGRATIONS = VECTOR_MIGRATIONS
     conn: C
     _deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]]
+    index_config: Optional[PostgresIndexConfig]
 
     def _get_batch_GET_ops_queries(
         self,
@@ -114,10 +191,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 +224,182 @@ 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)),
                     ]
                 )
+
+            # Then handle embeddings if configured
+            if self.index_config:
+                for op in inserts:
+                    if op.index is False:
+                        continue
+                    value = op.value
+                    ns = _namespace_to_text(op.namespace)
+                    k = op.key
+
+                    if op.index is None:
+                        paths = self.index_config["__tokenized_fields"]
+                    else:
+                        paths = [(ix, tokenize_path(ix)) for ix in op.index]
+
+                    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 = """
-                SELECT prefix, key, value, created_at, updated_at
-                FROM store
-                WHERE prefix LIKE %s
-            """
-            params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
+    ) -> 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):
+            # Build filter conditions first
+            filter_params = []
+            filter_conditions = []
             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)
+                            filter_params.extend(filter_params_)
                     else:
                         filter_conditions.append("value->%s = %s::jsonb")
-                        params.extend([key, json.dumps(value)])
-                query += " AND " + " AND ".join(filter_conditions)
+                        filter_params.extend([key, json.dumps(value)])
 
-            # Note: we will need to not do this if sim/keyword search
-            # is used
-            query += " ORDER BY updated_at DESC LIMIT %s OFFSET %s"
-            params.extend([op.limit, op.offset])
+            # Vector search branch
+            if op.query and self.index_config:
+                embedding_requests.append((idx, op.query))
 
-            queries.append((query, params))
-        return queries
+                score_operator = _get_distance_operator(self)
+                vector_type = (
+                    cast(PostgresIndexConfig, self.index_config)
+                    .get("ann_index_config", {})
+                    .get("vector_type", "vector")
+                )
+
+                if (
+                    vector_type == "bit"
+                    and self.index_config.get("distance_type") == "hamming"
+                ):
+                    score_operator = score_operator % (
+                        "%s",
+                        self.index_config["dims"],
+                    )
+                else:
+                    score_operator = score_operator % (
+                        "%s",
+                        vector_type,
+                    )
+
+                vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
+                expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
+
+                # Vector search with CTE for proper score handling
+                filter_str = (
+                    ""
+                    if not filter_conditions
+                    else " AND " + " AND ".join(filter_conditions)
+                )
+                base_query = f"""
+                    WITH scored AS (
+                        SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS score
+                        FROM store s
+                        JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
+                        WHERE s.prefix LIKE %s {filter_str}
+                        ORDER BY {score_operator} DESC 
+                        LIMIT %s
+                    )
+                    SELECT * FROM (
+                        SELECT DISTINCT ON (prefix, key) 
+                            prefix, key, value, created_at, updated_at, score 
+                        FROM scored 
+                        ORDER BY prefix, key, score DESC
+                    ) AS unique_docs
+                    ORDER BY score DESC
+                    LIMIT %s
+                    OFFSET %s
+                """
+                params = [
+                    _PLACEHOLDER,  # Vector placeholder
+                    f"{_namespace_to_text(op.namespace_prefix)}%",
+                    *filter_params,
+                    _PLACEHOLDER,
+                    expanded_limit,
+                    op.limit,
+                    op.offset,
+                ]
+
+            # Regular search branch
+            else:
+                base_query = """
+                    SELECT prefix, key, value, created_at, updated_at
+                    FROM store
+                    WHERE prefix LIKE %s
+                """
+                params = [f"{_namespace_to_text(op.namespace_prefix)}%"]
+
+                if filter_conditions:
+                    params.extend(filter_params)
+                    base_query += " AND " + " AND ".join(filter_conditions)
+
+                base_query += " ORDER BY updated_at DESC"
+                base_query += " LIMIT %s OFFSET %s"
+                params.extend([op.limit, op.offset])
+
+            queries.append((base_query, params))
+
+        return queries, embedding_requests
 
     def _get_batch_list_namespaces_queries(
         self,
@@ -249,13 +451,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 +491,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 +499,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 +513,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.
+            index (Optional[PostgresIndexConfig]): The index configuration for the store.
+
         Returns:
             PostgresStore: A new PostgresStore instance.
         """
@@ -310,16 +545,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 +654,32 @@ 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 +689,20 @@ 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):
+                _paramslist = queries[idx][1]
+                for i in range(len(_paramslist)):
+                    if _paramslist[i] is _PLACEHOLDER:
+                        _paramslist[i] = embedding
+
+        for (idx, _), (query, params) in zip(search_ops, queries):
+            # Execute the actual query
             cur.execute(query, params)
             rows = cast(list[Row], cur.fetchall())
             results[idx] = [
@@ -463,9 +734,10 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         already exist and runs database migrations. It MUST be called directly by the user
         the first time the store is used.
         """
-        with self._cursor() as cur:
+
+        def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
             try:
-                cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
+                cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
                 row = cast(dict, cur.fetchone())
                 if row is None:
                     version = -1
@@ -474,18 +746,35 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             except UndefinedTable:
                 version = -1
                 cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS store_migrations (
+                    f"""
+                    CREATE TABLE IF NOT EXISTS {table} (
                         v INTEGER PRIMARY KEY
                     )
                 """
                 )
-            for v, migration in enumerate(
-                self.MIGRATIONS[version + 1 :], start=version + 1
-            ):
-                cur.execute(migration)
+            return version
+
+        with self._cursor() as cur:
+            version = _get_version(cur, table="store_migrations")
+            for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
+                cur.execute(sql)
                 cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
 
+            if self.index_config:
+                version = _get_version(cur, table="vector_migrations")
+                for v, migration in enumerate(
+                    self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
+                ):
+                    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 vector_migrations (v) VALUES (%s)", (v,))
+
 
 class Row(TypedDict):
     key: str
@@ -495,6 +784,45 @@ class Row(TypedDict):
     updated_at: datetime
 
 
+# Private utilities
+
+_DEFAULT_ANN_CONFIG = ANNIndexConfig(
+    vector_type="vector",
+)
+
+
+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("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
+    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 _namespace_to_text(
     namespace: tuple[str, ...], handle_wildcards: bool = False
 ) -> str:
@@ -510,16 +838,26 @@ 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."""
-    loader = loader or _json_loads
+    """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"]
-    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"],
+    }
+
+    return Item(**kwargs)
 
 
 def _row_to_search_item(
@@ -575,3 +913,62 @@ 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."""
+    # Note: Today, we are not using ANN indices due to restrictions
+    # on PGVector's support for mixing vector and non-vector filters
+    # To use the index, PGVector expects:
+    #  - ORDER BY the operator NOT an expression (even negation blocks it)
+    #  - ASCENDING order
+    #  - Any WHERE clause should be over a partial index.
+    # If we violate any of these, it will use a sequential scan
+    # See https://github.com/pgvector/pgvector/issues/216 and the
+    # pgvector documentation for more details.
+    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
+
+
+_PLACEHOLDER = object()

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.4"
+version = "2.0.5"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
 
 [tool.poetry.dependencies]
 python = "^3.9.0,<4.0"
-langgraph-checkpoint = "^2.0.2"
+langgraph-checkpoint = "^2.0.7"
 orjson = ">=3.10.1"
 psycopg = "^3.0.0"
 psycopg-pool = "^3.0.0"

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,11 @@ 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)
+
+
+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 itertools
 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,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -181,272 +189,319 @@ 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(
+    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,
+        "ann_index_config": {
+            "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=[
+        (vector_type, distance_type)
+        for vector_type in VECTOR_TYPES
+        for distance_type in (
+            ["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
+        )
+    ],
+    ids=lambda p: f"{p[0]}_{p[1]}",
+)
+async def vector_store(
+    request,
+    fake_embeddings: CharacterEmbeddings,
+) -> AsyncIterator[AsyncPostgresStore]:
+    """Create a store with vector search enabled."""
+    vector_type, distance_type = request.param
+    async with _create_vector_store(
+        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(
+    "vector_type,distance_type",
+    [
+        *itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
+    ],
+)
+async def test_embed_with_path(
+    request: Any,
+    fake_embeddings: CharacterEmbeddings,
+    vector_type: str,
+    distance_type: str,
+) -> None:
+    """Test vector search with specific text fields in Postgres store."""
+    async with _create_vector_store(
+        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
+
+
+@pytest.mark.parametrize(
+    "vector_type,distance_type",
+    [
+        *itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
+    ],
+)
+async def test_search_sorting(
+    request: Any,
+    fake_embeddings: CharacterEmbeddings,
+    vector_type: str,
+    distance_type: str,
+) -> None:
+    """Test operation-level field configuration for vector search."""
+    async with _create_vector_store(
+        vector_type,
+        distance_type,
+        fake_embeddings,
+        text_fields=["key1"],  # Default fields that won't match our test data
+    ) as store:
+        amatch = {
+            "key1": "mmm",
+        }
+
+        await store.aput(("test", "M"), "M", amatch)
+        N = 100
+        for i in range(N):
+            await store.aput(("test", "A"), f"A{i}", {"key1": "no"})
+        for i in range(N):
+            await store.aput(("test", "Z"), f"Z{i}", {"key1": "no"})
+
+        results = await store.asearch(("test",), query="mmm", limit=10)
+        assert len(results) == 10
+        assert len(set(r.key for r in results)) == 10
+        assert results[0].key == "M"
+        assert results[0].score > results[1].score

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,11 @@ from langgraph.store.base import (
     SearchOp,
 )
 from langgraph.store.postgres import PostgresStore
+from tests.conftest import (
+    DEFAULT_URI,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -340,3 +347,351 @@ class TestPostgresStore:
             # Cleanup
             for namespace, key, _ in test_data:
                 store.delete(namespace, key)
+
+
+@contextmanager
+def _create_vector_store(
+    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,
+        "ann_index_config": {
+            "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=[
+        (vector_type, distance_type)
+        for vector_type in VECTOR_TYPES
+        for distance_type in (
+            ["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
+        )
+    ],
+    ids=lambda p: f"{p[0]}_{p[1]}",
+)
+def vector_store(
+    request,
+    fake_embeddings: Embeddings,
+) -> PostgresStore:
+    """Create a store with vector search enabled."""
+    vector_type, distance_type = request.param
+    with _create_vector_store(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(
+    "vector_type,distance_type",
+    [
+        ("vector", "cosine"),
+        ("vector", "inner_product"),
+        ("halfvec", "cosine"),
+        ("halfvec", "inner_product"),
+    ],
+)
+def test_embed_with_path_sync(
+    request: Any,
+    fake_embeddings: CharacterEmbeddings,
+    vector_type: str,
+    distance_type: str,
+) -> None:
+    """Test vector search with specific text fields in Postgres store."""
+    with _create_vector_store(
+        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
+
+
+@pytest.mark.parametrize(
+    "vector_type,distance_type",
+    [
+        ("vector", "cosine"),
+        ("vector", "inner_product"),
+        ("halfvec", "cosine"),
+        ("halfvec", "inner_product"),
+    ],
+)
+def test_embed_with_path_operation_config(
+    request: Any,
+    fake_embeddings: CharacterEmbeddings,
+    vector_type: str,
+    distance_type: str,
+) -> None:
+    """Test operation-level field configuration for vector search."""
+    with _create_vector_store(
+        vector_type,
+        distance_type,
+        fake_embeddings,
+        text_fields=["key17"],  # Default fields that won't match our test data
+    ) as store:
+        doc3 = {
+            "key0": "aaa",
+            "key1": "bbb",
+            "key2": "ccc",
+            "key3": "ddd",
+        }
+        doc4 = {
+            "key0": "eee",
+            "key1": "bbb",  # Same as doc3.key1
+            "key2": "fff",
+            "key3": "ggg",
+        }
+
+        store.put(("test",), "doc3", doc3, index=["key0", "key1"])
+        store.put(("test",), "doc4", doc4, index=["key1", "key3"])
+
+        results = store.search(("test",), query="aaa")
+        assert len(results) == 2
+        assert results[0].key == "doc3"
+        assert len(set(r.key for r in results)) == 2
+        assert results[0].score > results[1].score
+
+        results = store.search(("test",), query="ggg")
+        assert len(results) == 2
+        assert results[0].key == "doc4"
+        assert results[0].score > results[1].score
+
+        results = store.search(("test",), query="bbb")
+        assert len(results) == 2
+        assert results[0].key != results[1].key
+        assert results[0].score == pytest.approx(results[1].score, abs=1e-3)
+
+        results = store.search(("test",), query="ccc")
+        assert len(results) == 2
+        assert all(
+            r.score < 0.9 for r in results
+        )  # Unindexed field should have low scores
+
+        # Test index=False behavior
+        doc5 = {
+            "key0": "hhh",
+            "key1": "iii",
+        }
+        store.put(("test",), "doc5", doc5, index=False)
+        results = store.search(("test",))
+        assert len(results) == 3
+        assert all(r.score is None for r in results)
+        assert any(r.key == "doc5" for r in results)
+
+        results = store.search(("test",), query="hhh")
+        # TODO: We don't currently fill in additional results if there are not enough
+        # returned during vector search.
+        # assert len(results) == 3
+        # doc5_result = next(r for r in results if r.key == "doc5")
+        # assert doc5_result.score is None

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:

libs/checkpoint/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint"
-version = "2.0.5"
+version = "2.0.7"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 authors = []
 license = "MIT"

libs/sdk-py/langgraph_sdk/client.py

@@ -18,6 +18,7 @@ from typing import (
     Dict,
     Iterator,
     List,
+    Literal,
     Optional,
     Sequence,
     Union,
@@ -59,7 +60,6 @@ from langgraph_sdk.schema import (
     ThreadStatus,
     ThreadUpdateStateResponse,
 )
-from langgraph_sdk.sse import EventSource
 
 logger = logging.getLogger(__name__)
 
@@ -293,7 +293,7 @@ class HttpClient:
                 else:
                     logger.error(f"Error from langgraph-api: {body}", exc_info=e)
                 raise e
-            async for event in EventSource(sse.response).aiter_sse():
+            async for event in sse.aiter_sse():
                 yield StreamPart(
                     event.event, orjson.loads(event.data) if event.data else None
                 )
@@ -1947,7 +1947,7 @@ class CronClient:
 
         Example Usage:
 
-            cron_run = await client.crons.create(
+            cron_run = client.crons.create(
                 assistant_id="agent",
                 schedule="27 15 * * *",
                 input={"messages": [{"role": "user", "content": "hello!"}]},
@@ -2071,7 +2071,12 @@ class StoreClient:
         self.http = http
 
     async def put_item(
-        self, namespace: Sequence[str], /, key: str, value: dict[str, Any]
+        self,
+        namespace: Sequence[str],
+        /,
+        key: str,
+        value: dict[str, Any],
+        index: Optional[Union[Literal[False], list[str]]] = None,
     ) -> None:
         """Store or update an item.
 
@@ -2079,6 +2084,7 @@ class StoreClient:
             namespace: A list of strings representing the namespace path.
             key: The unique identifier for the item within the namespace.
             value: A dictionary containing the item's data.
+            index: Controls search indexing - None (use defaults), False (disable), or list of field paths to index.
 
         Returns:
             None
@@ -2096,11 +2102,7 @@ class StoreClient:
                 raise ValueError(
                     f"Invalid namespace label '{label}'. Namespace labels cannot contain periods ('.')."
                 )
-        payload = {
-            "namespace": namespace,
-            "key": key,
-            "value": value,
-        }
+        payload = {"namespace": namespace, "key": key, "value": value, "index": index}
         await self.http.put("/store/items", json=payload)
 
     async def get_item(self, namespace: Sequence[str], /, key: str) -> Item:
@@ -2168,6 +2170,7 @@ class StoreClient:
         filter: Optional[dict[str, Any]] = None,
         limit: int = 10,
         offset: int = 0,
+        query: Optional[str] = None,
     ) -> SearchItemsResponse:
         """Search for items within a namespace prefix.
 
@@ -2176,6 +2179,7 @@ class StoreClient:
             filter: Optional dictionary of key-value pairs to filter results.
             limit: Maximum number of items to return (default is 10).
             offset: Number of items to skip before returning results (default is 0).
+            query: Optional query for natural language search.
 
         Returns:
             List[Item]: A list of items matching the search criteria.
@@ -2213,6 +2217,7 @@ class StoreClient:
             "filter": filter,
             "limit": limit,
             "offset": offset,
+            "query": query,
         }
 
         return await self.http.post("/store/items/search", json=_provided_vals(payload))
@@ -2427,7 +2432,7 @@ class SyncHttpClient:
                 else:
                     logger.error(f"Error from langgraph-api: {body}", exc_info=e)
                 raise e
-            for event in EventSource(sse.response).iter_sse():
+            for event in sse.iter_sse():
                 yield StreamPart(
                     event.event, orjson.loads(event.data) if event.data else None
                 )
@@ -4155,7 +4160,12 @@ class SyncStoreClient:
         self.http = http
 
     def put_item(
-        self, namespace: Sequence[str], /, key: str, value: dict[str, Any]
+        self,
+        namespace: Sequence[str],
+        /,
+        key: str,
+        value: dict[str, Any],
+        index: Optional[Union[Literal[False], list[str]]] = None,
     ) -> None:
         """Store or update an item.
 
@@ -4163,6 +4173,7 @@ class SyncStoreClient:
             namespace: A list of strings representing the namespace path.
             key: The unique identifier for the item within the namespace.
             value: A dictionary containing the item's data.
+            index: Controls search indexing - None (use defaults), False (disable), or list of field paths to index.
 
         Returns:
             None
@@ -4184,6 +4195,7 @@ class SyncStoreClient:
             "namespace": namespace,
             "key": key,
             "value": value,
+            "index": index,
         }
         self.http.put("/store/items", json=payload)
 
@@ -4251,6 +4263,7 @@ class SyncStoreClient:
         filter: Optional[dict[str, Any]] = None,
         limit: int = 10,
         offset: int = 0,
+        query: Optional[str] = None,
     ) -> SearchItemsResponse:
         """Search for items within a namespace prefix.
 
@@ -4259,6 +4272,7 @@ class SyncStoreClient:
             filter: Optional dictionary of key-value pairs to filter results.
             limit: Maximum number of items to return (default is 10).
             offset: Number of items to skip before returning results (default is 0).
+            query: Optional query for natural language search.
 
         Returns:
             List[Item]: A list of items matching the search criteria.
@@ -4296,6 +4310,7 @@ class SyncStoreClient:
             "filter": filter,
             "limit": limit,
             "offset": offset,
+            "query": query,
         }
         return self.http.post("/store/items/search", json=_provided_vals(payload))
 

libs/sdk-py/langgraph_sdk/schema.py

@@ -325,10 +325,21 @@ class ListNamespaceResponse(TypedDict):
     """A list of namespace paths, where each path is a list of strings."""
 
 
+class SearchItem(Item, total=False):
+    """Item with an optional relevance score from search operations.
+
+    Attributes:
+        score (Optional[float]): Relevance/similarity score. Included when
+            searching a compatible store with a natural language query.
+    """
+
+    score: Optional[float]
+
+
 class SearchItemsResponse(TypedDict):
     """Response structure for searching items."""
 
-    items: list[Item]
+    items: list[SearchItem]
     """A list of items matching the search criteria."""
 
 

libs/sdk-py/langgraph_sdk/sse.py

@@ -1,106 +0,0 @@
-"""Adapted from httpx_sse to split lines on \n, \r, \r\n per the SSE spec."""
-
-import io
-from typing import AsyncIterator, Iterator
-
-import httpx
-import httpx_sse
-import httpx_sse._decoders
-
-
-class BytesLineDecoder:
-    """
-    Handles incrementally reading lines from text.
-
-    Has the same behaviour as the stdllib bytes splitlines,
-    but handling the input iteratively.
-    """
-
-    def __init__(self) -> None:
-        self.buffer = io.BytesIO()
-        self.trailing_cr: bool = False
-
-    def decode(self, text: bytes) -> list[bytes]:
-        # See https://docs.python.org/3/glossary.html#term-universal-newlines
-        NEWLINE_CHARS = b"\n\r"
-
-        # We always push a trailing `\r` into the next decode iteration.
-        if self.trailing_cr:
-            text = b"\r" + text
-            self.trailing_cr = False
-        if text.endswith(b"\r"):
-            self.trailing_cr = True
-            text = text[:-1]
-
-        if not text:
-            # NOTE: the edge case input of empty text doesn't occur in practice,
-            # because other httpx internals filter out this value
-            return []  # pragma: no cover
-
-        trailing_newline = text[-1] in NEWLINE_CHARS
-        lines = text.splitlines()
-
-        if len(lines) == 1 and not trailing_newline:
-            # No new lines, buffer the input and continue.
-            self.buffer.write(lines[0])
-            return []
-
-        if self.buffer:
-            # Include any existing buffer in the first portion of the
-            # splitlines result.
-            lines = [self.buffer.getvalue() + lines[0]] + lines[1:]
-            self.buffer.truncate(0)
-
-        if not trailing_newline:
-            # If the last segment of splitlines is not newline terminated,
-            # then drop it from our output and start a new buffer.
-            self.buffer.write(lines.pop())
-
-        return lines
-
-    def flush(self) -> list[bytes]:
-        if not self.buffer and not self.trailing_cr:
-            return []
-
-        lines = [self.buffer.getvalue()] if self.buffer else []
-        self.buffer.truncate(0)
-        self.trailing_cr = False
-        return lines
-
-
-async def aiter_lines_raw(response: httpx.Response) -> AsyncIterator[bytes]:
-    decoder = BytesLineDecoder()
-    async for chunk in response.aiter_bytes():
-        for line in decoder.decode(chunk):
-            yield line
-    for line in decoder.flush():
-        yield line
-
-
-def iter_lines_raw(response: httpx.Response) -> Iterator[bytes]:
-    decoder = BytesLineDecoder()
-    for chunk in response.iter_bytes():
-        for line in decoder.decode(chunk):
-            yield line
-    for line in decoder.flush():
-        yield line
-
-
-class EventSource(httpx_sse.EventSource):
-    async def aiter_sse(self) -> AsyncIterator[httpx_sse.ServerSentEvent]:
-        self._check_content_type()
-        decoder = httpx_sse._decoders.SSEDecoder()
-        async for line in aiter_lines_raw(self._response):
-            line = line.rstrip(b"\n")
-            sse = decoder.decode(line.decode())
-            if sse is not None:
-                yield sse
-
-    def iter_sse(self) -> Iterator[httpx_sse.ServerSentEvent]:
-        self._check_content_type()
-        decoder = httpx_sse._decoders.SSEDecoder()
-        for line in iter_lines_raw(self._response):
-            line = line.rstrip(b"\n")
-            sse = decoder.decode(line.decode())
-            if sse is not None:
-                yield sse

libs/sdk-py/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-sdk"
-version = "0.1.38"
+version = "0.1.40"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"