[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/langgraph/store/postgres/aio.py

@@ -21,6 +21,7 @@ from langgraph.store.base import (
 )
 from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.postgres.base import (
+    _PLACEHOLDER,
     BasePostgresStore,
     PoolConfig,
     PostgresIndexConfig,
@@ -147,11 +148,10 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         already exist and runs database migrations. It MUST be called directly by the user
         the first time the store is used.
         """
-        async with self._cursor() as cur:
+
+        async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
             try:
-                await cur.execute(
-                    "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
-                )
+                await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
                 row = await cur.fetchone()
                 if row is None:
                     version = -1
@@ -160,22 +160,25 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
             except UndefinedTable:
                 version = -1
                 await cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS store_migrations (
+                    f"""
+                    CREATE TABLE IF NOT EXISTS {table} (
                         v INTEGER PRIMARY KEY
                     )
                 """
                 )
+            return version
 
-            for v, migration in enumerate(
-                self.MIGRATIONS[version + 1 :], start=version + 1
-            ):
-                if isinstance(migration, str):
-                    sql = migration
-                else:
-                    if migration.condition and not migration.condition(self):
-                        continue
+        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 = {
@@ -183,9 +186,10 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                             for k, v in migration.params.items()
                         }
                         sql = sql % params
-
-                await cur.execute(sql)
-                await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+                    await cur.execute(sql)
+                    await cur.execute(
+                        "INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
+                    )
 
     async def _execute_batch(
         self,
@@ -289,7 +293,10 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                 [query for _, query in embedding_requests]
             )
             for (idx, _), vector in zip(embedding_requests, vectors):
-                queries[idx][1][0] = vector
+                _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)

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

@@ -55,16 +55,10 @@ class Migration(NamedTuple):
     """A database migration with optional conditions and parameters."""
 
     sql: str
-    condition: Optional[Callable[[Any], bool]] = None
     params: Optional[dict[str, Any]] = None
 
 
-def _embedding_requested(store: Any) -> bool:
-    """Check if vector operations are available in the database."""
-    return bool(store.index_config)
-
-
-MIGRATIONS: Sequence[Union[str, Migration]] = [
+MIGRATIONS: Sequence[str] = [
     """
 CREATE TABLE IF NOT EXISTS store (
     -- 'prefix' represents the doc's 'namespace'
@@ -80,11 +74,13 @@ CREATE TABLE IF NOT EXISTS store (
 -- For faster lookups by prefix
 CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
 """,
+]
+
+VECTOR_MIGRATIONS: Sequence[Migration] = [
     Migration(
         """
 CREATE EXTENSION IF NOT EXISTS vector;
 """,
-        condition=_embedding_requested,
     ),
     Migration(
         """
@@ -99,37 +95,20 @@ CREATE TABLE IF NOT EXISTS store_vectors (
     FOREIGN KEY (prefix, key) REFERENCES store(prefix, key) ON DELETE CASCADE
 );
 """,
-        condition=_embedding_requested,
         params={
             "dims": lambda store: store.index_config["dims"],
             "vector_type": lambda store: (
                 cast(PostgresIndexConfig, store.index_config)
-                .get("db_index_config", {})
+                .get("ann_index_config", {})
                 .get("vector_type", "vector")
             ),
         },
     ),
-    Migration(
-        """
-CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
-    USING %(index_type)s (embedding %(ops)s)%(index_params)s;
-""",
-        condition=_embedding_requested,
-        params={
-            "index_type": lambda store: _get_index_params(store)[0],
-            "ops": lambda store: _get_vector_type_ops(store),
-            "index_params": lambda store: (
-                " WITH ("
-                + ", ".join(f"{k}={v}" for k, v in _get_index_params(store)[1].items())
-                + ")"
-                if _get_index_params(store)[1]
-                else ""
-            ),
-        },
-    ),
+    # 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])
 
 
@@ -158,11 +137,9 @@ class PoolConfig(TypedDict, total=False):
     """
 
 
-class DBIndexConfig(TypedDict, total=False):
+class ANNIndexConfig(TypedDict, total=False):
     """Configuration for vector index in PostgreSQL store."""
 
-    kind: Literal["hnsw", "ivfflat"]
-    """Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat."""
     vector_type: Literal["vector", "halfvec"]
     """Type of vector storage to use.
     Options:
@@ -171,42 +148,13 @@ class DBIndexConfig(TypedDict, total=False):
     """
 
 
-class HNSWConfig(DBIndexConfig, total=False):
-    """Configuration for HNSW (Hierarchical Navigable Small World) index."""
-
-    kind: Literal["hnsw"]  # type: ignore[misc]
-    m: int
-    """Maximum number of connections per layer. Default is 16."""
-    ef_construction: int
-    """Size of dynamic candidate list for index construction. Default is 64."""
-
-
-class IVFFlatConfig(DBIndexConfig, total=False):
-    """IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
-
-    Three keys to achieving good recall are:
-    1. Create the index after the table has some data
-    2. Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for over 1M rows
-    3. When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) - a good place to start is sqrt(lists)
-    """
-
-    kind: Literal["ivfflat"]  # type: ignore[misc]
-    nlist: int
-    """Number of inverted lists (clusters) for IVF index.
-    
-    Determines the number of clusters used in the index structure.
-    Higher values can improve search speed but increase index size and build time.
-    Typically set to the square root of the number of vectors in the index.
-    """
-
-
 class PostgresIndexConfig(IndexConfig, total=False):
     """Configuration for vector embeddings in PostgreSQL store with pgvector-specific options.
 
     Extends EmbeddingConfig with additional configuration for pgvector index and vector types.
     """
 
-    db_index_config: Union[HNSWConfig, IVFFlatConfig]
+    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:
@@ -218,17 +166,11 @@ class PostgresIndexConfig(IndexConfig, total=False):
 
 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]
 
-    @staticmethod
-    def _get_default_index_config() -> IndexConfig:
-        return HNSWConfig(
-            kind="hnsw",
-            vector_type="vector",
-        )
-
     def _get_batch_GET_ops_queries(
         self,
         get_ops: Sequence[tuple[int, GetOp]],
@@ -298,13 +240,12 @@ class BasePostgresStore(Generic[C]):
                     [
                         _namespace_to_text(op.namespace),
                         op.key,
-                        Jsonb(cast(dict, op.value).copy()),
+                        Jsonb(cast(dict, op.value)),
                     ]
                 )
 
             # Then handle embeddings if configured
             if self.index_config:
-                paths = self.index_config["__tokenized_fields"]
                 for op in inserts:
                     if op.index is False:
                         continue
@@ -312,6 +253,11 @@ class BasePostgresStore(Generic[C]):
                     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):
@@ -355,22 +301,30 @@ class BasePostgresStore(Generic[C]):
         embedding_requests = []
 
         for idx, (_, op) in enumerate(search_ops):
-            base_query = """
-                SELECT prefix, key, value, created_at, updated_at
-                FROM store
-                WHERE prefix LIKE %s
-            """
-            params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
-            needs_vector_search = False
+            # Build filter conditions first
+            filter_params = []
+            filter_conditions = []
+            if op.filter:
+                for key, value in op.filter.items():
+                    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")
+                        filter_params.extend([key, json.dumps(value)])
 
+            # Vector search branch
             if op.query and self.index_config:
-                needs_vector_search = True
                 embedding_requests.append((idx, op.query))
 
-                score_expr = _get_distance_operator(self)
+                score_operator = _get_distance_operator(self)
                 vector_type = (
                     cast(PostgresIndexConfig, self.index_config)
-                    .get("db_index_config", self._get_default_index_config())
+                    .get("ann_index_config", {})
                     .get("vector_type", "vector")
                 )
 
@@ -378,61 +332,71 @@ class BasePostgresStore(Generic[C]):
                     vector_type == "bit"
                     and self.index_config.get("distance_type") == "hamming"
                 ):
-                    score_expr = score_expr % ("%s", self.index_config["dims"])
+                    score_operator = score_operator % (
+                        "%s",
+                        self.index_config["dims"],
+                    )
                 else:
-                    score_expr = score_expr % ("%s", vector_type)
+                    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
 
-                # Direct query with DISTINCT ON to get best score per document
+                # 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 DISTINCT ON (s.prefix, s.key)
-                            s.prefix, s.key, s.value, s.created_at, s.updated_at,
-                            {score_expr} as score
+                    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
-                        ORDER BY s.prefix, s.key, score DESC
+                        WHERE s.prefix LIKE %s {filter_str}
+                        ORDER BY {score_operator} DESC 
                         LIMIT %s
                     )
-
-                    SELECT * FROM scored
+                    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 = [
-                    None,  # Vector placeholder
+                    _PLACEHOLDER,  # Vector placeholder
                     f"{_namespace_to_text(op.namespace_prefix)}%",
+                    *filter_params,
+                    _PLACEHOLDER,
                     expanded_limit,
+                    op.limit,
+                    op.offset,
                 ]
 
-            if op.filter:
-                filter_conditions = []
-                for key, value in op.filter.items():
-                    if isinstance(value, dict):
-                        for op_name, val in value.items():
-                            condition, filter_params = self._get_filter_condition(
-                                key, op_name, val
-                            )
-                            filter_conditions.append(condition)
-                            params.extend(filter_params)
-                    else:
-                        filter_conditions.append("value->%s = %s::jsonb")
-                        params.extend([key, json.dumps(value)])
+            # 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:
-                    if needs_vector_search:
-                        base_query += " WHERE " + " AND ".join(filter_conditions)
-                    else:
-                        base_query += " AND " + " AND ".join(filter_conditions)
+                    params.extend(filter_params)
+                    base_query += " AND " + " AND ".join(filter_conditions)
 
-            if needs_vector_search:
-                base_query += " ORDER BY score DESC"
-            else:
                 base_query += " ORDER BY updated_at DESC"
+                base_query += " LIMIT %s OFFSET %s"
+                params.extend([op.limit, op.offset])
 
-            base_query += " LIMIT %s OFFSET %s"
-            params.extend([op.limit, op.offset])
             queries.append((base_query, params))
 
         return queries, embedding_requests
@@ -559,7 +523,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             pool_config (Optional[PoolArgs]): Configuration for the connection pool.
                 If provided, will create a connection pool and use it instead of a single connection.
                 This overrides the `pipeline` argument.
-            embedding (Optional[PostgresIndexConfig]): The embedding config.
+            index (Optional[PostgresIndexConfig]): The index configuration for the store.
 
         Returns:
             PostgresStore: A new PostgresStore instance.
@@ -705,7 +669,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             vectors = self.embeddings.embed_documents(
                 [param[-1] for param in txt_params]
             )
-
             queries.append(
                 (
                     query,
@@ -733,9 +696,13 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                 [query for _, query in embedding_requests]
             )
             for (idx, _), embedding in zip(embedding_requests, embeddings):
-                queries[idx][1][0] = embedding
+                _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] = [
@@ -767,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
@@ -778,21 +746,25 @@ 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
-            ):
-                if isinstance(migration, str):
-                    sql = migration
-                else:
-                    if migration.condition and not migration.condition(self):
-                        continue
+            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 = {
@@ -800,8 +772,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                             for k, v in migration.params.items()
                         }
                         sql = sql % params
-                cur.execute(sql)
-                cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+                    cur.execute(sql)
+                    cur.execute("INSERT INTO vector_migrations (v) VALUES (%s)", (v,))
 
 
 class Row(TypedDict):
@@ -814,6 +786,10 @@ class Row(TypedDict):
 
 # 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."""
@@ -821,9 +797,7 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
         return "vector_cosine_ops"
 
     config = cast(PostgresIndexConfig, store.index_config)
-    index_config = config.get(
-        "db_index_config", BasePostgresStore._get_default_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(
@@ -849,19 +823,6 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str:
     return f"{type_prefix}_{distance_suffix}"
 
 
-def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]:
-    """Get the index type and configuration based on config."""
-    if not store.index_config:
-        return "hnsw", {}
-
-    config = cast(PostgresIndexConfig, store.index_config)
-    default_config = BasePostgresStore._get_default_index_config()
-    index_config = config.get("db_index_config", default_config).copy()
-    kind = index_config.pop("kind", "hnsw")
-    index_config.pop("vector_type", None)
-    return kind, index_config
-
-
 def _namespace_to_text(
     namespace: tuple[str, ...], handle_wildcards: bool = False
 ) -> str:
@@ -925,32 +886,6 @@ def _row_to_search_item(
     )
 
 
-def _row_to_search_item(
-    namespace: tuple[str, ...],
-    row: Row,
-    *,
-    loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
-) -> SearchItem:
-    """Convert a row from the database into an Item."""
-    loader = loader or _json_loads
-    val = row["value"]
-    score = row.get("score")
-    if score is not None:
-        try:
-            score = float(score)  # type: ignore[arg-type]
-        except ValueError:
-            logger.warning("Invalid score: %s", score)
-            score = None
-    return SearchItem(
-        value=val if isinstance(val, dict) else loader(val),
-        key=row["key"],
-        namespace=namespace,
-        created_at=row["created_at"],
-        updated_at=row["updated_at"],
-        score=score,
-    )
-
-
 def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
     grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
     tot = 0
@@ -982,6 +917,15 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
 
 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 "
@@ -1025,3 +969,6 @@ def _ensure_index_config(
         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/conftest.py

@@ -40,5 +40,4 @@ def fake_embeddings() -> CharacterEmbeddings:
     return CharacterEmbeddings(dims=500)
 
 
-INDEX_TYPES = ["hnsw", "ivfflat"]
 VECTOR_TYPES = ["vector", "halfvec"]

libs/checkpoint-postgres/tests/test_async_store.py

@@ -1,4 +1,5 @@
 # type: ignore
+import itertools
 import sys
 import uuid
 from collections.abc import AsyncIterator
@@ -13,7 +14,6 @@ from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
 from langgraph.store.postgres import AsyncPostgresStore
 from tests.conftest import (
     DEFAULT_URI,
-    INDEX_TYPES,
     VECTOR_TYPES,
     CharacterEmbeddings,
 )
@@ -191,7 +191,6 @@ async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
 
 @asynccontextmanager
 async def _create_vector_store(
-    index_type: str,
     vector_type: str,
     distance_type: str,
     fake_embeddings: CharacterEmbeddings,
@@ -215,8 +214,7 @@ async def _create_vector_store(
     index_config = {
         "dims": fake_embeddings.dims,
         "embed": fake_embeddings,
-        "db_index_config": {
-            "kind": index_type,
+        "ann_index_config": {
             "vector_type": vector_type,
         },
         "distance_type": distance_type,
@@ -244,25 +242,22 @@ async def _create_vector_store(
 @pytest.fixture(
     scope="function",
     params=[
-        (index_type, vector_type, distance_type)
-        for index_type in INDEX_TYPES
+        (vector_type, distance_type)
         for vector_type in VECTOR_TYPES
         for distance_type in (
-            (["hamming"] if index_type == "ivfflat" else ["hamming", "jaccard"])
-            if vector_type == "bit"
-            else ["l2", "inner_product", "cosine"]
+            ["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
         )
     ],
-    ids=lambda p: f"{p[0]}_{p[1]}_{p[2]}",
+    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."""
-    index_type, vector_type, distance_type = request.param
+    vector_type, distance_type = request.param
     async with _create_vector_store(
-        index_type, vector_type, distance_type, fake_embeddings
+        vector_type, distance_type, fake_embeddings
     ) as store:
         yield store
 
@@ -417,24 +412,19 @@ async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> Non
 
 
 @pytest.mark.parametrize(
-    "index_type,vector_type,distance_type",
+    "vector_type,distance_type",
     [
-        ("ivfflat", "vector", "cosine"),
-        ("hnsw", "vector", "cosine"),
-        ("hnsw", "halfvec", "cosine"),
-        ("hnsw", "halfvec", "inner_product"),
+        *itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
     ],
 )
 async def test_embed_with_path(
     request: Any,
     fake_embeddings: CharacterEmbeddings,
-    index_type: str,
     vector_type: str,
     distance_type: str,
 ) -> None:
     """Test vector search with specific text fields in Postgres store."""
     async with _create_vector_store(
-        index_type,
         vector_type,
         distance_type,
         fake_embeddings,
@@ -478,3 +468,40 @@ async def test_embed_with_path(
         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

@@ -19,7 +19,6 @@ from langgraph.store.base import (
 from langgraph.store.postgres import PostgresStore
 from tests.conftest import (
     DEFAULT_URI,
-    INDEX_TYPES,
     VECTOR_TYPES,
     CharacterEmbeddings,
 )
@@ -352,7 +351,6 @@ class TestPostgresStore:
 
 @contextmanager
 def _create_vector_store(
-    index_type: str,
     vector_type: str,
     distance_type: str,
     fake_embeddings: Embeddings,
@@ -373,8 +371,7 @@ def _create_vector_store(
     index_config = {
         "dims": fake_embeddings.dims,
         "embed": fake_embeddings,
-        "db_index_config": {
-            "kind": index_type,
+        "ann_index_config": {
             "vector_type": vector_type,
         },
         "distance_type": distance_type,
@@ -398,26 +395,21 @@ def _create_vector_store(
 @pytest.fixture(
     scope="function",
     params=[
-        (index_type, vector_type, distance_type)
-        for index_type in INDEX_TYPES
+        (vector_type, distance_type)
         for vector_type in VECTOR_TYPES
         for distance_type in (
-            (["hamming"] if index_type == "ivfflat" else ["hamming", "jaccard"])
-            if vector_type == "bit"
-            else ["l2", "inner_product", "cosine"]
+            ["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
         )
     ],
-    ids=lambda p: f"{p[0]}_{p[1]}_{p[2]}",
+    ids=lambda p: f"{p[0]}_{p[1]}",
 )
 def vector_store(
     request,
     fake_embeddings: Embeddings,
 ) -> PostgresStore:
     """Create a store with vector search enabled."""
-    index_type, vector_type, distance_type = request.param
-    with _create_vector_store(
-        index_type, vector_type, distance_type, fake_embeddings
-    ) as store:
+    vector_type, distance_type = request.param
+    with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
         yield store
 
 
@@ -555,24 +547,22 @@ def test_vector_search_edge_cases(vector_store: PostgresStore) -> None:
 
 
 @pytest.mark.parametrize(
-    "index_type,vector_type,distance_type",
+    "vector_type,distance_type",
     [
-        ("ivfflat", "vector", "cosine"),
-        ("hnsw", "vector", "cosine"),
-        ("hnsw", "halfvec", "cosine"),
-        ("hnsw", "halfvec", "inner_product"),
+        ("vector", "cosine"),
+        ("vector", "inner_product"),
+        ("halfvec", "cosine"),
+        ("halfvec", "inner_product"),
     ],
 )
 def test_embed_with_path_sync(
     request: Any,
     fake_embeddings: CharacterEmbeddings,
-    index_type: str,
     vector_type: str,
     distance_type: str,
 ) -> None:
     """Test vector search with specific text fields in Postgres store."""
     with _create_vector_store(
-        index_type,
         vector_type,
         distance_type,
         fake_embeddings,
@@ -626,3 +616,82 @@ def test_embed_with_path_sync(
         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/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.append(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.37"
+version = "0.1.40"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"