Test put-time fields

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

@@ -35,7 +35,6 @@ from langgraph.store.base import (
     ListNamespacesOp,
     Op,
     PutOp,
-    ResponseMetadata,
     Result,
     SearchItem,
     SearchOp,
@@ -532,20 +531,20 @@ def _row_to_search_item(
     """Convert a row from the database into an Item."""
     loader = loader or _json_loads
     val = row["value"]
-    response_metadata: Optional[ResponseMetadata] = (
-        {
-            "score": float(row["score"]),
-        }
-        if row.get("score") is not None
-        else None
-    )
+    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"],
-        response_metadata=response_metadata,
+        score=score,
     )
 
 

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

@@ -152,35 +152,80 @@ class SearchOp(NamedTuple):
 
 
 class PutOp(NamedTuple):
-    """Operation to store, update, or delete an item."""
+    """Operation to store, update, or delete an item in the store.
+
+    This class represents a single operation to modify the store's contents,
+    whether adding new items, updating existing ones, or removing them.
+    """
 
     namespace: tuple[str, ...]
-    """Hierarchical path for the item.
-    
-    Represented as a tuple of strings, allowing for nested categorization.
-    For example: ("documents", "user123")
+    """Hierarchical path that identifies the location of the item.
+
+    The namespace acts as a folder-like structure to organize items.
+    Each element in the tuple represents one level in the hierarchy.
+
+    Examples:
+        ("documents",) - Root level documents
+        ("documents", "user123") - User-specific documents
+        ("cache", "embeddings", "v1") - Nested cache structure
     """
 
     key: str
-    """Unique identifier for the document.
-    
-    Should be distinct within its namespace.
+    """Unique identifier for the item within its namespace.
+
+    The key must be unique within the specific namespace to avoid conflicts.
+    Together with the namespace, it forms a complete path to the item.
+
+    Example:
+        If namespace is ("documents", "user123") and key is "report1",
+        the full path would effectively be "documents/user123/report1"
     """
 
     value: Optional[dict[str, Any]]
-    """Data to be stored, or None to delete the item.
-    
+    """The data to store, or None to mark the item for deletion.
+
+    The value must be a dictionary with string keys and JSON-serializable values.
+    Setting this to None signals that the item should be deleted.
+
     Schema:
-    - Should be a dictionary where:
-      - Keys are strings representing field names
-      - Values can be of any serializable type
-    - If None, it indicates that the item should be deleted
+        {
+            "field1": "string value",
+            "field2": 123,
+            "nested": {"can": "contain", "any": "serializable data"}
+        }
     """
-    index: Optional[bool] = None  # type: ignore[assignment]
-    """Whether to index the item (if supported by the store).
-    
-    Defaults to True if the store supports indexing. This will embed the document
-    so it can be queried using search.
+
+    index: Optional[Union[Literal[False], list[str]]] = None  # type: ignore[assignment]
+    """Controls how the item's fields are indexed for search operations.
+
+    Indexing configuration determines how the item can be found through search:
+    - None (default): Uses the store's default indexing configuration (if provided)
+    - False: Disables indexing for this item
+    - list[str]: Specifies which json path fields to index for search
+
+    The item remains accessible through direct get() operations regardless of indexing.
+    When indexed, fields can be searched using natural language queries through
+    vector similarity search (if supported by the store implementation).
+
+    Path Syntax:
+        - Simple field access: "field"
+        - Nested fields: "parent.child.grandchild"
+        - Array indexing:
+          - Specific index: "array[0]"
+          - Last element: "array[-1]"
+          - All elements (each individually): "array[*]"
+
+    Examples:
+        None - Use store defaults
+        False - Don't index this item
+        [
+            "metadata.title",                    # Nested field access
+            "chapters[*].content",               # Index content from all chapters as separate vectors
+            "authors[0].name",                   # First author's name
+            "revisions[-1].changes",             # Most recent revision's changes
+            "sections[*].paragraphs[*].text",    # All text from all paragraphs in all sections
+            "metadata.tags[*]",                  # All tags in metadata
+        ]
     """
 
 
@@ -320,16 +365,43 @@ class BaseStore(ABC):
         namespace: tuple[str, ...],
         key: str,
         value: dict[str, Any],
-        index: Optional[bool] = None,
+        index: Optional[Union[Literal[False], list[str]]] = None,
     ) -> None:
-        """Store or update an item.
+        """Store or update an item in the store.
 
         Args:
-            namespace: Hierarchical path for the item.
-            key: Unique identifier within the namespace.
-            value: Dictionary containing the item's data.
-            index: Whether to index the item (if supported by the store).
-                Defaults to True if the store supports indexing.
+            namespace: Hierarchical path for the item, represented as a tuple of strings.
+                Example: ("documents", "user123")
+            key: Unique identifier within the namespace. Together with namespace forms
+                the complete path to the item.
+            value: Dictionary containing the item's data. Must contain string keys
+                and JSON-serializable values.
+            index: Controls how the item's fields are indexed for search:
+                - None (default): Use store's default indexing configuration
+                - False: Disable indexing for this item
+                - list[str]: List of field paths to index, supporting:
+                    - Nested fields: "metadata.title"
+                    - Array access: "chapters[*].content" (each indexed separately)
+                    - Specific indices: "authors[0].name"
+
+        Note:
+            Indexing capabilities depend on your store implementation.
+            Some implementations may support only a subset of indexing features.
+
+        Examples:
+            # Simple storage without special indexing
+            store.put(("docs",), "report", {"title": "Annual Report"})
+
+            # Index specific fields for search
+            store.put(
+                ("docs",),
+                "report",
+                {
+                    "title": "Q4 Report",
+                    "chapters": [{"content": "..."}, {"content": "..."}]
+                },
+                index=["title", "chapters[*].content"]
+            )
         """
         _validate_namespace(namespace)
         self.batch([PutOp(namespace, key, value, index=index)])
@@ -439,19 +511,46 @@ class BaseStore(ABC):
         namespace: tuple[str, ...],
         key: str,
         value: dict[str, Any],
-        index: Optional[bool] = None,
+        index: Optional[Union[Literal[False], list[str]]] = None,
     ) -> None:
-        """Asynchronously store or update an item.
+        """Asynchronously store or update an item in the store.
 
         Args:
-            namespace: Hierarchical path for the item.
-            key: Unique identifier within the namespace.
-            value: Dictionary containing the item's data.
-            index: Whether to index the item (if supported by the store).
-                Defaults to True if the store supports indexing.
+            namespace: Hierarchical path for the item, represented as a tuple of strings.
+                Example: ("documents", "user123")
+            key: Unique identifier within the namespace. Together with namespace forms
+                the complete path to the item.
+            value: Dictionary containing the item's data. Must contain string keys
+                and JSON-serializable values.
+            index: Controls how the item's fields are indexed for search:
+                - None (default): Use store's default indexing configuration
+                - False: Disable indexing for this item
+                - list[str]: List of field paths to index, supporting:
+                    - Nested fields: "metadata.title"
+                    - Array access: "chapters[*].content" (each indexed separately)
+                    - Specific indices: "authors[0].name"
+
+        Note:
+            Indexing capabilities depend on your store implementation.
+            Some implementations may support only a subset of indexing features.
+
+        Examples:
+            # Simple storage without special indexing
+            await store.aput(("docs",), "report", {"title": "Annual Report"})
+
+            # Index specific fields for search
+            await store.aput(
+                ("docs",),
+                "report",
+                {
+                    "title": "Q4 Report",
+                    "chapters": [{"content": "..."}, {"content": "..."}]
+                },
+                index=["title", "chapters[*].content"]
+            )
         """
         _validate_namespace(namespace)
-        await self.abatch([PutOp(namespace, key, value, index)])
+        await self.abatch([PutOp(namespace, key, value, index=index)])
 
     async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
         """Asynchronously delete an item.

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

@@ -1,6 +1,6 @@
 import asyncio
 import weakref
-from typing import Any, Optional
+from typing import Any, Literal, Optional, Union
 
 from langgraph.store.base import (
     BaseStore,
@@ -58,7 +58,7 @@ class AsyncBatchedBaseStore(BaseStore):
         namespace: tuple[str, ...],
         key: str,
         value: dict[str, Any],
-        index: Optional[bool] = None,
+        index: Optional[Union[Literal[False], list[str]]] = None,
     ) -> None:
         _validate_namespace(namespace)
         fut = self._loop.create_future()

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

@@ -63,34 +63,41 @@ class EmbeddingsLambda(Embeddings):
     """Wrapper to convert embedding functions into LangChain's Embeddings interface.
 
     This class allows arbitrary embedding functions to be used with LangChain-compatible
-    tools. It supports both synchronous and asynchronous operations, and can be
-    initialized with either:
-    1. A synchronous function for both sync/async operations
-    2. An async function for both sync/async operations
-    3. Both sync and async functions for their respective operations
+    tools. It supports both synchronous and asynchronous operations, and can handle:
+    1. A synchronous function for sync operations (async operations will use sync function)
+    2. An async function for both sync/async operations (sync operations will raise an error)
 
     The embedding functions should convert text into fixed-dimensional vectors that
     capture the semantic meaning of the text.
 
     Args:
         func: Function that converts text to embeddings. Can be sync or async.
-            If async, it will be used for both sync and async operations.
-        afunc: Optional async function for embeddings. If provided, it will be used
-            for async operations while func is used for sync operations.
-            Must be None if func is async.
+            If async, it will be used for async operations, but sync operations
+            will raise an error. If sync, it will be used for both sync and async operations.
 
     Example:
+        >>> # With a sync function
         >>> def my_embed_fn(texts):
         ...     # Return 2D embeddings for each text
         ...     return [[0.1, 0.2] for _ in texts]
         >>> embeddings = EmbeddingsLambda(my_embed_fn)
         >>> result = embeddings.embed_query("hello")  # Returns [0.1, 0.2]
+        >>> await embeddings.aembed_query("hello")  # Also returns [0.1, 0.2]
+        >>>
+        >>> # With an async function
+        >>> async def my_async_fn(texts):
+        ...     return [[0.1, 0.2] for _ in texts]
+        >>> embeddings = EmbeddingsLambda(my_async_fn)
+        >>> await embeddings.aembed_query("hello")  # Returns [0.1, 0.2]
+        >>> # Note: embed_query() would raise an error
     """
 
     def __init__(
         self,
-        func: Union[EmbeddingsFunc, AEmbeddingsFunc, None],
+        func: Union[EmbeddingsFunc, AEmbeddingsFunc],
     ) -> None:
+        if func is None:
+            raise ValueError("func must be provided")
         if _is_async_callable(func):
             self.afunc = func
         else:

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

@@ -233,17 +233,21 @@ class InMemoryStore(BaseStore):
             if op.query and queryinmem_store:
                 query_embedding = queryinmem_store[op.query]
                 flat_items, flat_vectors = [], []
+                scoreless = []
                 for item, vectors in candidates:
                     for vector in vectors:
                         flat_items.append(item)
                         flat_vectors.append(vector)
+                    if not vectors:
+                        scoreless.append(item)
+
                 scores = _cosine_similarity(query_embedding, flat_vectors)
                 sorted_results = sorted(
                     zip(scores, flat_items), key=lambda x: x[0], reverse=True
                 )
                 # max pooling
                 seen: set[tuple[tuple[str, ...], str]] = set()
-                kept = []
+                kept: list[tuple[Optional[float], Item]] = []
                 for score, item in sorted_results:
                     key = (item.namespace, item.key)
                     if key in seen:
@@ -256,6 +260,12 @@ class InMemoryStore(BaseStore):
                         continue
 
                     kept.append((score, item))
+                if scoreless and len(kept) < op.limit:
+                    # Corner case: if we request more items than what we have embedded,
+                    # fill the rest with non-scored items
+                    kept.extend(
+                        (None, item) for item in scoreless[: op.limit - len(kept)]
+                    )
 
                 results[i] = [
                     SearchItem(
@@ -264,7 +274,7 @@ class InMemoryStore(BaseStore):
                         value=item.value,
                         created_at=item.created_at,
                         updated_at=item.updated_at,
-                        score=float(score),
+                        score=float(score) if score is not None else None,
                     )
                     for score, item in kept
                 ]
@@ -331,7 +341,11 @@ class InMemoryStore(BaseStore):
 
             for op in put_ops.values():
                 if op.value is not None and op.index is not False:
-                    for path, field in self.index_config["__tokenized_fields"]:
+                    if op.index is None:
+                        paths = self.index_config["__tokenized_fields"]
+                    else:
+                        paths = [(ix, tokenize_path(ix)) for ix in op.index]
+                    for path, field in paths:
                         texts = get_text_at_path(op.value, field)
                         if texts:
                             if len(texts) > 1:

libs/checkpoint/tests/test_store.py

@@ -426,6 +426,9 @@ async def test_cannot_put_empty_namespace() -> None:
     assert val is not None
     assert val.value == doc
     assert (await async_store.asearch(("foo", "langgraph", "foo")))[0].value == doc
+    assert (await async_store.asearch(("foo", "langgraph", "foo"), query="bar"))[
+        0
+    ].value == doc
     await async_store.adelete(("foo", "langgraph", "foo"), "bar")
     assert (await async_store.aget(("foo", "langgraph", "foo"), "bar")) is None
 
@@ -840,54 +843,8 @@ async def test_async_vector_search_pagination(
     assert len(all_results) == 5
 
 
-def test_vector_search_edge_cases(fake_embeddings: CharacterEmbeddings) -> None:
-    """Test edge cases in vector search."""
-    store = InMemoryStore(
-        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
-    )
-    store.put(("test",), "doc1", {"text": "test document"})
-
-    results = store.search(("test",), query="")
-    assert len(results) == 1
-
-    results = store.search(("test",), query=None)
-    assert len(results) == 1
-
-    long_query = "test " * 100
-    results = store.search(("test",), query=long_query)
-    assert len(results) == 1
-
-    special_query = "test!@#$%^&*()"
-    results = store.search(("test",), query=special_query)
-    assert len(results) == 1
-
-
-async def test_async_vector_search_edge_cases(
-    fake_embeddings: CharacterEmbeddings,
-) -> None:
-    """Test edge cases in vector search using async methods."""
-    store = InMemoryStore(
-        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
-    )
-    await store.aput(("test",), "doc1", {"text": "test document"})
-
-    results = await store.asearch(("test",), query="")
-    assert len(results) == 1
-
-    results = await store.asearch(("test",), query=None)
-    assert len(results) == 1
-
-    long_query = "test " * 100
-    results = await store.asearch(("test",), query=long_query)
-    assert len(results) == 1
-
-    special_query = "test!@#$%^&*()"
-    results = await store.asearch(("test",), query=special_query)
-    assert len(results) == 1
-
-
 async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None:
-    # Basi
+    # Test store-level field configuration
     store = InMemoryStore(
         index={
             "dims": fake_embeddings.dims,
@@ -935,3 +892,58 @@ async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None:
     assert len(results) == 2
     assert results[0].score < ascore
     assert results[1].score < ascore
+
+    # Test operation-level field configuration
+    store_no_defaults = InMemoryStore(
+        index={
+            "dims": fake_embeddings.dims,
+            "embed": fake_embeddings,
+            "fields": ["key17"],
+        }
+    )
+
+    doc3 = {
+        "key0": "aaa",
+        "key1": "bbb",
+        "key2": "ccc",
+        "key3": "ddd",
+    }
+    doc4 = {
+        "key0": "eee",
+        "key1": "bbb",  # Same as doc3.key1
+        "key2": "fff",
+        "key3": "ggg",
+    }
+
+    await store_no_defaults.aput(("test",), "doc3", doc3, index=["key0", "key1"])
+    await store_no_defaults.aput(("test",), "doc4", doc4, index=["key1", "key3"])
+
+    results = await store_no_defaults.asearch(("test",), query="aaa")
+    assert len(results) == 2
+    assert results[0].key == "doc3"
+    assert results[0].score is not None and results[0].score > results[1].score
+
+    results = await store_no_defaults.asearch(("test",), query="ggg")
+    assert len(results) == 2
+    assert results[0].key == "doc4"
+    assert results[0].score is not None and results[0].score > results[1].score
+
+    results = await store_no_defaults.asearch(("test",), query="bbb")
+    assert len(results) == 2
+    assert results[0].key != results[1].key
+    assert results[0].score == results[1].score
+
+    results = await store_no_defaults.asearch(("test",), query="ccc")
+    assert len(results) == 2
+    assert all(r.score < ascore for r in results)
+
+    doc5 = {
+        "key0": "hhh",
+        "key1": "iii",
+    }
+    await store_no_defaults.aput(("test",), "doc5", doc5, index=False)
+
+    results = await store_no_defaults.asearch(("test",), query="hhh")
+    assert len(results) == 3
+    doc5_result = next(r for r in results if r.key == "doc5")
+    assert doc5_result.score is None