From 0ed26a3af9a2e20554da853d4f57a6453ed036b5 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 27 Nov 2024 00:23:29 -0800
Subject: [PATCH 1/5] handle no emb situation

---
 libs/checkpoint/langgraph/store/memory/__init__.py | 2 +-
 libs/checkpoint/tests/test_store.py                | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py
index b350c8a17..7af33962a 100644
--- a/libs/checkpoint/langgraph/store/memory/__init__.py
+++ b/libs/checkpoint/langgraph/store/memory/__init__.py
@@ -232,7 +232,7 @@ class InMemoryStore(BaseStore):
             if not candidates:
                 results[i] = []
                 continue
-            if op.query:
+            if op.query and queryinmem_store:
                 query_embedding = queryinmem_store[op.query]
                 flat_items, flat_vectors = [], []
                 for item, vectors in candidates:
diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py
index 8c5df5b04..d00fc21bf 100644
--- a/libs/checkpoint/tests/test_store.py
+++ b/libs/checkpoint/tests/test_store.py
@@ -382,7 +382,9 @@ async def test_cannot_put_empty_namespace() -> None:
 
     await store.aput(("foo", "langgraph", "foo"), "bar", doc)
     assert (await store.aget(("foo", "langgraph", "foo"), "bar")).value == doc  # type: ignore[union-attr]
-    assert (await store.asearch(("foo", "langgraph", "foo")))[0].value == doc
+    assert (await store.asearch(("foo", "langgraph", "foo"), query="bar"))[
+        0
+    ].value == doc
     await store.adelete(("foo", "langgraph", "foo"), "bar")
     assert (await store.aget(("foo", "langgraph", "foo"), "bar")) is None
     store.put(("foo", "langgraph", "foo"), "bar", doc)

From d75492c330fe4e98e4413855c904cbff7af7d79f Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 27 Nov 2024 11:20:03 -0800
Subject: [PATCH 2/5] Rename config

---
 .../langgraph/store/base/__init__.py          | 10 ++--------
 .../checkpoint/langgraph/store/base/_embed.py | 20 +++----------------
 .../langgraph/store/memory/__init__.py        |  4 ++--
 3 files changed, 7 insertions(+), 27 deletions(-)

diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py
index ef6729d83..f1353d052 100644
--- a/libs/checkpoint/langgraph/store/base/__init__.py
+++ b/libs/checkpoint/langgraph/store/base/__init__.py
@@ -227,8 +227,8 @@ class InvalidNamespaceError(ValueError):
     """Provided namespace is invalid."""
 
 
-class EmbeddingConfig(TypedDict, total=False):
-    """Configuration for vector embeddings in PostgreSQL store."""
+class IndexConfig(TypedDict, total=False):
+    """Configuration for indexing documents for semantic search in the store."""
 
     dims: int
     """Number of dimensions in the embedding vectors.
@@ -245,12 +245,6 @@ class EmbeddingConfig(TypedDict, total=False):
 
     embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc]
     """Optional function to generate embeddings from text."""
-    aembed: Optional[AEmbeddingsFunc]
-    """Optional asynchronous function to generate embeddings from text.
-    
-    Provide for asynchronous embedding generation if you do not provide
-    an Embeddings object.
-    """
 
     text_fields: Optional[list[str]]
     """Fields to extract text from for embedding generation.
diff --git a/libs/checkpoint/langgraph/store/base/_embed.py b/libs/checkpoint/langgraph/store/base/_embed.py
index a04b4f138..a6c88e1f6 100644
--- a/libs/checkpoint/langgraph/store/base/_embed.py
+++ b/libs/checkpoint/langgraph/store/base/_embed.py
@@ -29,8 +29,6 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi
 
 def ensure_embeddings(
     embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None],
-    *,
-    aembed: Optional[AEmbeddingsFunc] = None,
 ) -> Embeddings:
     """Ensure that an embedding function conforms to LangChain's Embeddings interface.
 
@@ -42,9 +40,6 @@ def ensure_embeddings(
         embed: Either an existing Embeddings instance, or a function that converts
             text to embeddings. If the function is async, it will be used for both
             sync and async operations.
-        aembed: Optional async function for embeddings. If provided, it will be used
-            for async operations while the sync function is used for sync operations.
-            Must be None if embed is async.
 
     Returns:
         An Embeddings instance that wraps the provided function(s).
@@ -56,14 +51,12 @@ def ensure_embeddings(
         >>> embeddings = ensure_embeddings(my_embed_fn)
         >>> # Wrap an async function
         >>> embeddings = ensure_embeddings(my_async_fn)
-        >>> # Provide both sync and async implementations
-        >>> embeddings = ensure_embeddings(my_embed_fn, aembed=my_async_fn)
     """
-    if embed is None and aembed is None:
-        raise ValueError("embed or aembed must be provided")
+    if embed is None:
+        raise ValueError("embed must be provided")
     if isinstance(embed, Embeddings):
         return embed
-    return EmbeddingsLambda(embed, afunc=aembed)
+    return EmbeddingsLambda(embed)
 
 
 class EmbeddingsLambda(Embeddings):
@@ -97,18 +90,11 @@ class EmbeddingsLambda(Embeddings):
     def __init__(
         self,
         func: Union[EmbeddingsFunc, AEmbeddingsFunc, None],
-        afunc: Optional[AEmbeddingsFunc] = None,
     ) -> None:
         if _is_async_callable(func):
-            if afunc is not None:
-                raise ValueError(
-                    "afunc must be None if func is async. The async func will be used for both sync and async operations."
-                )
             self.afunc = func
         else:
             self.func = func
-            if afunc is not None:
-                self.afunc = afunc
 
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
         """Embed a list of texts into vectors.
diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py
index 7af33962a..2d246eadd 100644
--- a/libs/checkpoint/langgraph/store/memory/__init__.py
+++ b/libs/checkpoint/langgraph/store/memory/__init__.py
@@ -41,7 +41,7 @@ from langchain_core.embeddings import Embeddings
 
 from langgraph.store.base import (
     BaseStore,
-    EmbeddingConfig,
+    IndexConfig,
     GetOp,
     Item,
     ListNamespacesOp,
@@ -101,7 +101,7 @@ class InMemoryStore(BaseStore):
         "_vectors",
     )
 
-    def __init__(self, embedding_config: Optional[EmbeddingConfig] = None) -> None:
+    def __init__(self, embedding_config: Optional[IndexConfig] = None) -> None:
         self._data: dict[tuple[str, ...], dict[str, Item]] = defaultdict(dict)
         # [ns][key][path]
         self.inmem_store: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (

From 34a4ca3eaf4d94492f76ff407a477781f2c3b95e Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 27 Nov 2024 11:52:11 -0800
Subject: [PATCH 3/5] Rename file

---
 .../langgraph/store/base/__init__.py          | 23 ++----
 .../store/base/{_embed.py => embed.py}        |  2 +-
 .../langgraph/store/memory/__init__.py        | 56 +++++++-------
 .../embed_test_utils.py}                      |  0
 libs/checkpoint/tests/test_store.py           | 74 ++++++++++---------
 5 files changed, 74 insertions(+), 81 deletions(-)
 rename libs/checkpoint/langgraph/store/base/{_embed.py => embed.py} (99%)
 rename libs/checkpoint/{langgraph/store/base/_embed_test_utils.py => tests/embed_test_utils.py} (100%)

diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py
index f1353d052..71abcf208 100644
--- a/libs/checkpoint/langgraph/store/base/__init__.py
+++ b/libs/checkpoint/langgraph/store/base/__init__.py
@@ -15,7 +15,7 @@ from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Unio
 
 from langchain_core.embeddings import Embeddings
 
-from langgraph.store.base._embed import (
+from langgraph.store.base.embed import (
     AEmbeddingsFunc,
     EmbeddingsFunc,
     ensure_embeddings,
@@ -88,17 +88,10 @@ class Item:
         }
 
 
-class ResponseMetadata(TypedDict, total=False):
-    """Additional metadata about the response/result."""
-
-    score: float
-    """Relevance/similarity score if from a ranked operation."""
-
-
 class SearchItem(Item):
     """Represents a result item with additional response metadata."""
 
-    __slots__ = "response_metadata"
+    __slots__ = ("score",)
 
     def __init__(
         self,
@@ -107,7 +100,7 @@ class SearchItem(Item):
         value: dict[str, Any],
         created_at: datetime,
         updated_at: datetime,
-        response_metadata: Optional[ResponseMetadata] = None,
+        score: Optional[float] = None,
     ) -> None:
         """Initialize a result item.
 
@@ -117,7 +110,7 @@ class SearchItem(Item):
             value: The stored value.
             created_at: When the item was first created.
             updated_at: When the item was last updated.
-            response_metadata: Optional metadata about the response/result.
+            score: Relevance/similarity score if from a ranked operation.
         """
         super().__init__(
             value=value,
@@ -126,11 +119,11 @@ class SearchItem(Item):
             created_at=created_at,
             updated_at=updated_at,
         )
-        self.response_metadata = response_metadata or {}
+        self.score = score
 
     def dict(self) -> dict:
         result = super().dict()
-        result["response_metadata"] = self.response_metadata
+        result["score"] = self.score
         return result
 
 
@@ -246,10 +239,10 @@ class IndexConfig(TypedDict, total=False):
     embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc]
     """Optional function to generate embeddings from text."""
 
-    text_fields: Optional[list[str]]
+    fields: Optional[list[str]]
     """Fields to extract text from for embedding generation.
     
-    Defaults to ["__root__"], which embeds the json object as a whole.
+    Defaults to the root ["$"], which embeds the json object as a whole.
     """
 
 
diff --git a/libs/checkpoint/langgraph/store/base/_embed.py b/libs/checkpoint/langgraph/store/base/embed.py
similarity index 99%
rename from libs/checkpoint/langgraph/store/base/_embed.py
rename to libs/checkpoint/langgraph/store/base/embed.py
index a6c88e1f6..b8fc43814 100644
--- a/libs/checkpoint/langgraph/store/base/_embed.py
+++ b/libs/checkpoint/langgraph/store/base/embed.py
@@ -179,7 +179,7 @@ def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]:
             - Multi-field selection: "{field1,field2}"
             - Nested paths in multi-field: "{field1,nested.field2}"
     """
-    if not path or path == "__root__":
+    if not path or path == "$":
         return [json.dumps(obj, sort_keys=True)]
 
     tokens = tokenize_path(path) if isinstance(path, str) else path
diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py
index 2d246eadd..5f2768b13 100644
--- a/libs/checkpoint/langgraph/store/memory/__init__.py
+++ b/libs/checkpoint/langgraph/store/memory/__init__.py
@@ -11,7 +11,7 @@ Examples:
 
     Vector search with embeddings:
         from langchain_openai import OpenAIEmbeddings
-        store = InMemoryStore(embedding_config={
+        store = InMemoryStore(index={
             "dims": 1536,
             "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
         })
@@ -41,8 +41,8 @@ from langchain_core.embeddings import Embeddings
 
 from langgraph.store.base import (
     BaseStore,
-    IndexConfig,
     GetOp,
+    IndexConfig,
     Item,
     ListNamespacesOp,
     MatchCondition,
@@ -70,7 +70,7 @@ class InMemoryStore(BaseStore):
 
         Vector search with embeddings:
             from langchain_openai import OpenAIEmbeddings
-            store = InMemoryStore(embedding_config={
+            store = InMemoryStore(index={
                 "dims": 1536,
                 "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
             })
@@ -95,32 +95,32 @@ class InMemoryStore(BaseStore):
 
     __slots__ = (
         "_data",
-        "embedding_config",
-        "inmem_store",
-        "embeddings",
         "_vectors",
+        "index_config",
+        "embeddings",
     )
 
-    def __init__(self, embedding_config: Optional[IndexConfig] = None) -> None:
+    def __init__(self, *, index: Optional[IndexConfig] = None) -> None:
+        # Both _data and _vectors are wrapped in the In-memory API
+        # Do not change their names
         self._data: dict[tuple[str, ...], dict[str, Item]] = defaultdict(dict)
         # [ns][key][path]
-        self.inmem_store: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (
+        self._vectors: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (
             defaultdict(lambda: defaultdict(dict))
         )
-        self.embedding_config = embedding_config
-        if self.embedding_config:
-            self.embedding_config = self.embedding_config.copy()
+        self.index_config = index
+        if self.index_config:
+            self.index_config = self.index_config.copy()
             self.embeddings: Optional[Embeddings] = ensure_embeddings(
-                self.embedding_config.get("embed"),
-                aembed=self.embedding_config.get("aembed"),
+                self.index_config.get("embed"),
             )
-            self.embedding_config["__tokenized_fields"] = [
-                (p, tokenize_path(p)) if p != "__root__" else (p, p)
-                for p in (self.embedding_config.get("text_fields") or ["__root__"])
+            self.index_config["__tokenized_fields"] = [
+                (p, tokenize_path(p)) if p != "$" else (p, p)
+                for p in (self.index_config.get("fields") or ["$"])
             ]
 
         else:
-            self.embedding_config = None
+            self.index_config = None
             self.embeddings = None
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:
@@ -132,7 +132,7 @@ class InMemoryStore(BaseStore):
             self._batch_search(search_ops, queryinmem_store, results)
 
         to_embed = self._extract_texts(put_ops)
-        if to_embed and self.embedding_config and self.embeddings:
+        if to_embed and self.index_config and self.embeddings:
             embeddings = self.embeddings.embed_documents(list(to_embed))
             self._insertinmem_store(to_embed, embeddings)
         self._apply_put_ops(put_ops)
@@ -147,7 +147,7 @@ class InMemoryStore(BaseStore):
             self._batch_search(search_ops, queryinmem_store, results)
 
         to_embed = self._extract_texts(put_ops)
-        if to_embed and self.embedding_config and self.embeddings:
+        if to_embed and self.index_config and self.embeddings:
             embeddings = await self.embeddings.aembed_documents(list(to_embed))
             self._insertinmem_store(to_embed, embeddings)
         self._apply_put_ops(put_ops)
@@ -179,9 +179,7 @@ class InMemoryStore(BaseStore):
 
             for key, item in self._data[namespace].items():
                 if filter_func(item):
-                    if op.query and (
-                        embeddings := self.inmem_store[namespace].get(key)
-                    ):
+                    if op.query and (embeddings := self._vectors[namespace].get(key)):
                         filtered.append((item, list(embeddings.values())))
                     else:
                         filtered.append((item, []))
@@ -192,7 +190,7 @@ class InMemoryStore(BaseStore):
         search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
     ) -> dict[str, list[float]]:
         queryinmem_store = {}
-        if self.embedding_config and self.embeddings and search_ops:
+        if self.index_config and self.embeddings and search_ops:
             queries = {op.query for (op, _) in search_ops.values() if op.query}
 
             if queries:
@@ -211,7 +209,7 @@ class InMemoryStore(BaseStore):
         search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
     ) -> dict[str, list[float]]:
         queryinmem_store = {}
-        if self.embedding_config and self.embeddings and search_ops:
+        if self.index_config and self.embeddings and search_ops:
             queries = {op.query for (op, _) in search_ops.values() if op.query}
 
             if queries:
@@ -266,7 +264,7 @@ class InMemoryStore(BaseStore):
                         value=item.value,
                         created_at=item.created_at,
                         updated_at=item.updated_at,
-                        response_metadata={"score": float(score)},
+                        score=float(score),
                     )
                     for score, item in kept
                 ]
@@ -315,7 +313,7 @@ class InMemoryStore(BaseStore):
         for (namespace, key), op in put_ops.items():
             if op.value is None:
                 self._data[namespace].pop(key, None)
-                self.inmem_store[namespace].pop(key, None)
+                self._vectors[namespace].pop(key, None)
             else:
                 self._data[namespace][key] = Item(
                     value=op.value,
@@ -328,12 +326,12 @@ class InMemoryStore(BaseStore):
     def _extract_texts(
         self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]
     ) -> dict[str, list[tuple[tuple[str, ...], str, str]]]:
-        if put_ops and self.embedding_config and self.embeddings:
+        if put_ops and self.index_config and self.embeddings:
             to_embed = defaultdict(list)
 
             for op in put_ops.values():
                 if op.value is not None and op.index is not False:
-                    for path, field in self.embedding_config["__tokenized_fields"]:
+                    for path, field in self.index_config["__tokenized_fields"]:
                         texts = get_text_at_path(op.value, field)
                         if texts:
                             if len(texts) > 1:
@@ -361,7 +359,7 @@ class InMemoryStore(BaseStore):
                 f" match number of indices ({len(indices)})"
             )
         for embedding, (ns, key, path) in zip(embeddings, indices):
-            self.inmem_store[ns][key][path] = embedding
+            self._vectors[ns][key][path] = embedding
 
     def _handle_list_namespaces(self, op: ListNamespacesOp) -> list[tuple[str, ...]]:
         all_namespaces = list(
diff --git a/libs/checkpoint/langgraph/store/base/_embed_test_utils.py b/libs/checkpoint/tests/embed_test_utils.py
similarity index 100%
rename from libs/checkpoint/langgraph/store/base/_embed_test_utils.py
rename to libs/checkpoint/tests/embed_test_utils.py
diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py
index d00fc21bf..5df8da5d1 100644
--- a/libs/checkpoint/tests/test_store.py
+++ b/libs/checkpoint/tests/test_store.py
@@ -1,3 +1,4 @@
+# mypy: disable-error-code="operator"
 import asyncio
 import json
 from datetime import datetime
@@ -15,9 +16,9 @@ from langgraph.store.base import (
     Result,
     get_text_at_path,
 )
-from langgraph.store.base._embed_test_utils import CharacterEmbeddings
 from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.memory import InMemoryStore
+from tests.embed_test_utils import CharacterEmbeddings
 
 
 class MockAsyncBatchedStore(AsyncBatchedBaseStore):
@@ -51,7 +52,7 @@ def test_get_text_at_path() -> None:
         "empty_dict": {},
     }
 
-    assert get_text_at_path(nested_data, "__root__") == [
+    assert get_text_at_path(nested_data, "$") == [
         json.dumps(nested_data, sort_keys=True)
     ]
 
@@ -389,7 +390,7 @@ async def test_cannot_put_empty_namespace() -> None:
     assert (await store.aget(("foo", "langgraph", "foo"), "bar")) is None
     store.put(("foo", "langgraph", "foo"), "bar", doc)
     assert store.get(("foo", "langgraph", "foo"), "bar").value == doc  # type: ignore[union-attr]
-    assert store.search(("foo", "langgraph", "foo"))[0].value == doc
+    assert store.search(("foo", "langgraph", "foo"), query="bar")[0].value == doc
     store.delete(("foo", "langgraph", "foo"), "bar")
     assert store.get(("foo", "langgraph", "foo"), "bar") is None
 
@@ -510,11 +511,11 @@ def fake_embeddings() -> CharacterEmbeddings:
 def test_vector_store_initialization(fake_embeddings: CharacterEmbeddings) -> None:
     """Test store initialization with embedding config."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
-    assert store.embedding_config is not None
-    assert store.embedding_config["dims"] == fake_embeddings.dims
-    assert store.embedding_config["embed"] == fake_embeddings
+    assert store.index_config is not None
+    assert store.index_config["dims"] == fake_embeddings.dims
+    assert store.index_config["embed"] == fake_embeddings
 
 
 def test_vector_insert_with_auto_embedding(
@@ -522,7 +523,7 @@ def test_vector_insert_with_auto_embedding(
 ) -> None:
     """Test inserting items that get auto-embedded."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     docs = [
         ("doc1", {"text": "short text"}),
@@ -549,7 +550,7 @@ async def test_async_vector_insert_with_auto_embedding(
 ) -> None:
     """Test inserting items that get auto-embedded using async methods."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     docs = [
         ("doc1", {"text": "short text"}),
@@ -574,7 +575,7 @@ async def test_async_vector_insert_with_auto_embedding(
 def test_vector_update_with_embedding(fake_embeddings: CharacterEmbeddings) -> None:
     """Test that updating items properly updates their embeddings."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
     store.put(("test",), "doc2", {"text": "something about dogs"})
@@ -583,20 +584,20 @@ def test_vector_update_with_embedding(fake_embeddings: CharacterEmbeddings) -> N
     results_initial = store.search(("test",), query="Zany Xerxes")
     assert len(results_initial) > 0
     assert results_initial[0].key == "doc1"
-    initial_score = results_initial[0].response_metadata["score"]
+    initial_score = results_initial[0].score
+    assert initial_score is not None
 
     store.put(("test",), "doc1", {"text": "new text about dogs"})
 
     results_after = store.search(("test",), query="Zany Xerxes")
-    after_score = next(
-        (r.response_metadata["score"] for r in results_after if r.key == "doc1"), 0.0
-    )
+    after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
+    assert after_score is not None
     assert after_score < initial_score
 
     results_new = store.search(("test",), query="new text about dogs")
     for r in results_new:
         if r.key == "doc1":
-            assert r.response_metadata["score"] > after_score
+            assert r.score > after_score
 
     # Don't index this one
     store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
@@ -609,7 +610,7 @@ async def test_async_vector_update_with_embedding(
 ) -> None:
     """Test that updating items properly updates their embeddings using async methods."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     await store.aput(("test",), "doc1", {"text": "zany zebra Xerxes"})
     await store.aput(("test",), "doc2", {"text": "something about dogs"})
@@ -618,20 +619,20 @@ async def test_async_vector_update_with_embedding(
     results_initial = await store.asearch(("test",), query="Zany Xerxes")
     assert len(results_initial) > 0
     assert results_initial[0].key == "doc1"
-    initial_score = results_initial[0].response_metadata["score"]
+    initial_score = results_initial[0].score
 
     await store.aput(("test",), "doc1", {"text": "new text about dogs"})
 
     results_after = await store.asearch(("test",), query="Zany Xerxes")
-    after_score = next(
-        (r.response_metadata["score"] for r in results_after if r.key == "doc1"), 0.0
-    )
+    after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
+    assert after_score is not None
     assert after_score < initial_score
 
     results_new = await store.asearch(("test",), query="new text about dogs")
     for r in results_new:
         if r.key == "doc1":
-            assert r.response_metadata["score"] > after_score
+            assert r.score is not None
+            assert r.score > after_score
 
     # Don't index this one
     await store.aput(("test",), "doc4", {"text": "new text about dogs"}, index=False)
@@ -642,7 +643,7 @@ async def test_async_vector_update_with_embedding(
 def test_vector_search_with_filters(fake_embeddings: CharacterEmbeddings) -> None:
     """Test combining vector search with filters."""
     inmem_store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     # Insert test documents
     docs = [
@@ -682,7 +683,7 @@ async def test_async_vector_search_with_filters(
 ) -> None:
     """Test combining vector search with filters using async methods."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     # Insert test documents
     docs = [
@@ -722,7 +723,7 @@ async def test_async_batched_vector_search_concurrent(
 ) -> None:
     """Test concurrent vector search operations using async batched store."""
     store = MockAsyncBatchedStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
 
     colors = ["red", "blue", "green", "yellow", "purple"]
@@ -802,7 +803,7 @@ async def test_async_batched_vector_search_concurrent(
 def test_vector_search_pagination(fake_embeddings: CharacterEmbeddings) -> None:
     """Test pagination with vector search."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     for i in range(5):
         store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
@@ -823,7 +824,7 @@ async def test_async_vector_search_pagination(
 ) -> None:
     """Test pagination with vector search using async methods."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     for i in range(5):
         await store.aput(("test",), f"doc{i}", {"text": f"test document number {i}"})
@@ -842,7 +843,7 @@ async def test_async_vector_search_pagination(
 def test_vector_search_edge_cases(fake_embeddings: CharacterEmbeddings) -> None:
     """Test edge cases in vector search."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     store.put(("test",), "doc1", {"text": "test document"})
 
@@ -866,7 +867,7 @@ async def test_async_vector_search_edge_cases(
 ) -> None:
     """Test edge cases in vector search using async methods."""
     store = InMemoryStore(
-        embedding_config={"dims": fake_embeddings.dims, "embed": fake_embeddings}
+        index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
     )
     await store.aput(("test",), "doc1", {"text": "test document"})
 
@@ -888,11 +889,11 @@ async def test_async_vector_search_edge_cases(
 async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None:
     # Basi
     store = InMemoryStore(
-        embedding_config={
+        index={
             "dims": fake_embeddings.dims,
             "embed": fake_embeddings,
             # Key 2 isn't included. Don't index it.
-            "text_fields": ["key0", "key1", "key3"],
+            "fields": ["key0", "key1", "key3"],
         }
     )
     # This will have 2 vectors representing it
@@ -916,20 +917,21 @@ async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None:
     results = await store.asearch(("test",), query="xxx")
     assert len(results) == 2
     assert results[0].key != results[1].key
-    ascore = results[0].response_metadata["score"]
-    bscore = results[1].response_metadata["score"]
+    ascore = results[0].score
+    bscore = results[1].score
     assert ascore == bscore
+    assert ascore is not None and bscore is not 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].response_metadata["score"] > results[1].response_metadata["score"]
-    assert ascore == pytest.approx(results[0].response_metadata["score"], abs=1e-5)
+    assert results[0].score is not None and results[0].score > results[1].score
+    assert ascore == pytest.approx(results[0].score, abs=1e-5)
 
     # 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].response_metadata["score"] < ascore
-    assert results[1].response_metadata["score"] < ascore
+    assert results[0].score < ascore
+    assert results[1].score < ascore

From 112a4f6c12e31a5a7fe20e9787f4dd440c7301b0 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 27 Nov 2024 13:04:59 -0800
Subject: [PATCH 4/5] Test put-time fields

---
 .../langgraph/store/postgres/base.py          |  17 +-
 .../langgraph/store/base/__init__.py          | 167 ++++++++++++++----
 libs/checkpoint/langgraph/store/base/batch.py |   4 +-
 libs/checkpoint/langgraph/store/base/embed.py |  27 +--
 .../langgraph/store/memory/__init__.py        |  20 ++-
 libs/checkpoint/tests/test_store.py           | 106 ++++++-----
 6 files changed, 236 insertions(+), 105 deletions(-)

diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py
index 0dfce8871..200218c0d 100644
--- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py
+++ b/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,
     )
 
 
diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py
index 71abcf208..a7c0d25b0 100644
--- a/libs/checkpoint/langgraph/store/base/__init__.py
+++ b/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.
diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py
index bb74eae8e..e6b00efdb 100644
--- a/libs/checkpoint/langgraph/store/base/batch.py
+++ b/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()
diff --git a/libs/checkpoint/langgraph/store/base/embed.py b/libs/checkpoint/langgraph/store/base/embed.py
index b8fc43814..9d869b6ef 100644
--- a/libs/checkpoint/langgraph/store/base/embed.py
+++ b/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:
diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py
index 5f2768b13..40011a7ba 100644
--- a/libs/checkpoint/langgraph/store/memory/__init__.py
+++ b/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:
diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py
index 5df8da5d1..053d9a981 100644
--- a/libs/checkpoint/tests/test_store.py
+++ b/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

From c2a1f3a30a23b18ff2c682dd68f9fab87b998794 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 27 Nov 2024 14:06:54 -0800
Subject: [PATCH 5/5] Improve docs

---
 .../langgraph/store/base/__init__.py          | 386 ++++++++++++++----
 libs/checkpoint/langgraph/store/base/batch.py |   6 +-
 libs/checkpoint/langgraph/store/base/embed.py |  87 ++--
 3 files changed, 358 insertions(+), 121 deletions(-)

diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py
index a7c0d25b0..033cfba12 100644
--- a/libs/checkpoint/langgraph/store/base/__init__.py
+++ b/libs/checkpoint/langgraph/store/base/__init__.py
@@ -128,27 +128,248 @@ class SearchItem(Item):
 
 
 class GetOp(NamedTuple):
-    """Operation to retrieve an item by namespace and key."""
+    """Operation to retrieve a specific item by its namespace and key.
+
+    This operation allows precise retrieval of stored items using their full path
+    (namespace) and unique identifier (key) combination.
+
+    ??? example "Examples"
+
+        Basic item retrieval:
+        ```python
+        GetOp(namespace=("users", "profiles"), key="user123")
+        GetOp(namespace=("cache", "embeddings"), key="doc456")
+        ```
+    """
 
     namespace: tuple[str, ...]
-    """Hierarchical path for the item."""
+    """Hierarchical path that uniquely identifies the item's location.
+
+    ??? example "Examples"
+
+        ```python
+        ("users",)  # Root level users namespace
+        ("users", "profiles")  # Profiles within users namespace
+        ```
+    """
+
     key: str
-    """Unique identifier within the namespace."""
+    """Unique identifier for the item within its specific namespace.
+
+    ??? example "Examples"
+
+        ```python
+        "user123"  # For a user profile
+        "doc456"  # For a document
+        ```
+    """
 
 
 class SearchOp(NamedTuple):
-    """Operation to search for items within a namespace prefix."""
+    """Operation to search for items within a specified namespace hierarchy.
+
+    This operation supports both structured filtering and natural language search
+    within a given namespace prefix. It provides pagination through limit and offset
+    parameters.
+
+    Note:
+        Natural language search support depends on your store implementation.
+
+    ??? example "Examples"
+        Search with filters and pagination:
+        ```python
+        SearchOp(
+            namespace_prefix=("documents",),
+            filter={"type": "report", "status": "active"},
+            limit=5,
+            offset=10
+        )
+        ```
+
+        Natural language search:
+        ```python
+        SearchOp(
+            namespace_prefix=("users", "content"),
+            query="technical documentation about APIs",
+            limit=20
+        )
+        ```
+    """
 
     namespace_prefix: tuple[str, ...]
-    """Hierarchical path prefix to search within."""
+    """Hierarchical path prefix defining the search scope.
+
+    ??? example "Examples"
+
+        ```python
+        ()  # Search entire store
+        ("documents",)  # Search all documents
+        ("users", "content")  # Search within user content
+        ```
+    """
+
     filter: Optional[dict[str, Any]] = None
-    """Key-value pairs to filter results."""
+    """Key-value pairs for filtering results based on exact matches or comparison operators.
+
+    The filter supports both exact matches and operator-based comparisons.
+
+    Supported Operators:
+        - $eq: Equal to (same as direct value comparison)
+        - $ne: Not equal to
+        - $gt: Greater than
+        - $gte: Greater than or equal to
+        - $lt: Less than
+        - $lte: Less than or equal to
+
+    ??? example "Examples"
+
+        Simple exact match:
+
+        ```python
+        {"status": "active"}
+        ```
+
+        Comparison operators:
+
+        ```python
+        {"score": {"$gt": 4.99}}  # Score greater than 4.99
+        ```
+
+        Multiple conditions:
+
+        ```python
+        {
+            "score": {"$gte": 3.0},
+            "color": "red"
+        }
+        ```
+
+    Note:
+        Comparison operator support depends on your store implementation.
+    """
+
     limit: int = 10
-    """Maximum number of items to return."""
+    """Maximum number of items to return in the search results."""
+
     offset: int = 0
-    """Number of items to skip before returning results."""
+    """Number of matching items to skip for pagination."""
+
     query: Optional[str] = None
-    """The search query for natural language search."""
+    """Natural language search query for semantic search capabilities.
+
+    ??? example "Examples"
+        - "technical documentation about REST APIs"
+        - "machine learning papers from 2023"
+    """
+
+
+# Type representing a namespace path that can include wildcards
+NamespacePath = tuple[Union[str, Literal["*"]], ...]
+"""A tuple representing a namespace path that can include wildcards.
+
+Examples:
+    ("users",)  # Exact users namespace
+    ("documents", "*")  # Any sub-namespace under documents
+    ("cache", "*", "v1")  # Any cache category with v1 version
+"""
+
+# Type for specifying how to match namespaces
+NamespaceMatchType = Literal["prefix", "suffix"]
+"""Specifies how to match namespace paths.
+
+Values:
+    "prefix": Match from the start of the namespace
+    "suffix": Match from the end of the namespace
+"""
+
+
+class MatchCondition(NamedTuple):
+    """Represents a pattern for matching namespaces in the store.
+
+    This class combines a match type (prefix or suffix) with a namespace path
+    pattern that can include wildcards to flexibly match different namespace
+    hierarchies.
+
+    ??? example "Examples"
+        Prefix matching:
+        ```python
+        MatchCondition(match_type="prefix", path=("users", "profiles"))
+        ```
+
+        Suffix matching with wildcard:
+        ```python
+        MatchCondition(match_type="suffix", path=("cache", "*"))
+        ```
+
+        Simple suffix matching:
+        ```python
+        MatchCondition(match_type="suffix", path=("v1",))
+        ```
+    """
+
+    match_type: NamespaceMatchType
+    """Type of namespace matching to perform."""
+
+    path: NamespacePath
+    """Namespace path pattern that can include wildcards."""
+
+
+class ListNamespacesOp(NamedTuple):
+    """Operation to list and filter namespaces in the store.
+
+    This operation allows exploring the organization of data, finding specific
+    collections, and navigating the namespace hierarchy.
+
+    ??? example "Examples"
+
+        List all namespaces under the "documents" path:
+        ```python
+        ListNamespacesOp(
+            match_conditions=(MatchCondition(match_type="prefix", path=("documents",)),),
+            max_depth=2
+        )
+        ```
+
+        List all namespaces that end with "v1":
+        ```python
+        ListNamespacesOp(
+            match_conditions=(MatchCondition(match_type="suffix", path=("v1",)),),
+            limit=50
+        )
+        ```
+
+    """
+
+    match_conditions: Optional[tuple[MatchCondition, ...]] = None
+    """Optional conditions for filtering namespaces.
+
+    ??? example "Examples"
+        All user namespaces:
+        ```python
+        (MatchCondition(match_type="prefix", path=("users",)),)
+        ```
+
+        All namespaces that start with "docs" and end with "draft":
+        ```python
+        (
+            MatchCondition(match_type="prefix", path=("docs",)),
+            MatchCondition(match_type="suffix", path=("draft",))
+        ) 
+        ```
+    """
+
+    max_depth: Optional[int] = None
+    """Maximum depth of namespace hierarchy to return.
+
+    Note:
+        Namespaces deeper than this level will be truncated.
+    """
+
+    limit: int = 100
+    """Maximum number of namespaces to return."""
+
+    offset: int = 0
+    """Number of namespaces to skip for pagination."""
 
 
 class PutOp(NamedTuple):
@@ -164,10 +385,21 @@ class PutOp(NamedTuple):
     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
+    ??? example "Examples"
+        Root level documents
+        ```python
+        ("documents",)
+        ```
+        
+        User-specific documents
+        ```python
+        ("documents", "user123")
+        ```
+        
+        Nested cache structure
+        ```python
+        ("cache", "embeddings", "v1")
+        ```
     """
 
     key: str
@@ -187,7 +419,7 @@ class PutOp(NamedTuple):
     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:
+    Example:
         {
             "field1": "string value",
             "field2": 123,
@@ -215,9 +447,12 @@ class PutOp(NamedTuple):
           - Last element: "array[-1]"
           - All elements (each individually): "array[*]"
 
-    Examples:
-        None - Use store defaults
-        False - Don't index this item
+    ??? example "Examples"
+        - None - Use store defaults
+        - False - Don't index this item
+        - list[str] - List of fields to index
+        
+        ```python
         [
             "metadata.title",                    # Nested field access
             "chapters[*].content",               # Index content from all chapters as separate vectors
@@ -226,37 +461,10 @@ class PutOp(NamedTuple):
             "sections[*].paragraphs[*].text",    # All text from all paragraphs in all sections
             "metadata.tags[*]",                  # All tags in metadata
         ]
+        ```
     """
 
 
-NameSpacePath = tuple[Union[str, Literal["*"]], ...]
-
-NamespaceMatchType = Literal["prefix", "suffix"]
-
-
-class MatchCondition(NamedTuple):
-    """Represents a single match condition."""
-
-    match_type: NamespaceMatchType
-    path: NameSpacePath
-
-
-class ListNamespacesOp(NamedTuple):
-    """Operation to list namespaces with optional match conditions."""
-
-    match_conditions: Optional[tuple[MatchCondition, ...]] = None
-    """A tuple of match conditions to apply to namespaces."""
-
-    max_depth: Optional[int] = None
-    """Return namespaces up to this depth in the hierarchy."""
-
-    limit: int = 100
-    """Maximum number of namespaces to return."""
-
-    offset: int = 0
-    """Number of namespaces to skip before returning results."""
-
-
 Op = Union[GetOp, SearchOp, PutOp, ListNamespacesOp]
 Result = Union[Item, list[Item], list[SearchItem], list[tuple[str, ...]], None]
 
@@ -388,20 +596,21 @@ class BaseStore(ABC):
             Indexing capabilities depend on your store implementation.
             Some implementations may support only a subset of indexing features.
 
-        Examples:
-            # Simple storage without special indexing
+        ??? example "Examples"
+            Simple storage without special indexing (respects store defaults)
+            ```python
             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"]
-            )
+            Index specific fields for search
+            ```python
+            store.put(("docs",), "report", {"title": "Annual Report"}, index=["title"])
+            ```
+
+            Do not index for semantic search
+            ```python
+            store.put(("docs",), "report", {"title": "Annual Report"}, index=False)
+            ```
         """
         _validate_namespace(namespace)
         self.batch([PutOp(namespace, key, value, index=index)])
@@ -418,8 +627,8 @@ class BaseStore(ABC):
     def list_namespaces(
         self,
         *,
-        prefix: Optional[NameSpacePath] = None,
-        suffix: Optional[NameSpacePath] = None,
+        prefix: Optional[NamespacePath] = None,
+        suffix: Optional[NamespacePath] = None,
         max_depth: Optional[int] = None,
         limit: int = 100,
         offset: int = 0,
@@ -433,7 +642,7 @@ class BaseStore(ABC):
             prefix (Optional[Tuple[str, ...]]): Filter namespaces that start with this path.
             suffix (Optional[Tuple[str, ...]]): Filter namespaces that end with this path.
             max_depth (Optional[int]): Return namespaces up to this depth in the hierarchy.
-                Namespaces deeper than this level will be truncated to this depth.
+                Namespaces deeper than this level will be truncated.
             limit (int): Maximum number of namespaces to return (default 100).
             offset (int): Number of namespaces to skip for pagination (default 0).
 
@@ -441,16 +650,18 @@ class BaseStore(ABC):
             List[Tuple[str, ...]]: A list of namespace tuples that match the criteria.
             Each tuple represents a full namespace path up to `max_depth`.
 
-        Examples:
-
+        ??? example "Examples":
             Setting max_depth=3. Given the namespaces:
-                # ("a", "b", "c")
-                # ("a", "b", "d", "e")
-                # ("a", "b", "d", "i")
-                # ("a", "b", "f")
-                # ("a", "c", "f")
-                store.list_namespaces(prefix=("a", "b"), max_depth=3)
-                # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
+            ```python
+            # Example if you have the following namespaces:
+            # ("a", "b", "c")
+            # ("a", "b", "d", "e")
+            # ("a", "b", "d", "i")
+            # ("a", "b", "f")
+            # ("a", "c", "f")
+            store.list_namespaces(prefix=("a", "b"), max_depth=3)
+            # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
+            ```
         """
         match_conditions = []
         if prefix:
@@ -534,11 +745,14 @@ class BaseStore(ABC):
             Indexing capabilities depend on your store implementation.
             Some implementations may support only a subset of indexing features.
 
-        Examples:
-            # Simple storage without special indexing
+        ??? example "Examples"
+            Simple storage without special indexing:
+            ```python
             await store.aput(("docs",), "report", {"title": "Annual Report"})
+            ```
 
-            # Index specific fields for search
+            Index specific fields for search:
+            ```python
             await store.aput(
                 ("docs",),
                 "report",
@@ -548,6 +762,7 @@ class BaseStore(ABC):
                 },
                 index=["title", "chapters[*].content"]
             )
+            ```
         """
         _validate_namespace(namespace)
         await self.abatch([PutOp(namespace, key, value, index=index)])
@@ -564,8 +779,8 @@ class BaseStore(ABC):
     async def alist_namespaces(
         self,
         *,
-        prefix: Optional[NameSpacePath] = None,
-        suffix: Optional[NameSpacePath] = None,
+        prefix: Optional[NamespacePath] = None,
+        suffix: Optional[NamespacePath] = None,
         max_depth: Optional[int] = None,
         limit: int = 100,
         offset: int = 0,
@@ -587,16 +802,19 @@ class BaseStore(ABC):
             List[Tuple[str, ...]]: A list of namespace tuples that match the criteria.
             Each tuple represents a full namespace path up to `max_depth`.
 
-        Examples:
+        ??? example "Examples"
+            Setting max_depth=3 with existing namespaces:
+            ```python
+            # Given the following namespaces:
+            # ("a", "b", "c")
+            # ("a", "b", "d", "e")
+            # ("a", "b", "d", "i")
+            # ("a", "b", "f")
+            # ("a", "c", "f")
 
-            Setting max_depth=3. Given the namespaces:
-                # ("a", "b", "c")
-                # ("a", "b", "d", "e")
-                # ("a", "b", "d", "i")
-                # ("a", "b", "f")
-                # ("a", "c", "f")
-                await store.alist_namespaces(prefix=("a", "b"), max_depth=3)
-                # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
+            await store.alist_namespaces(prefix=("a", "b"), max_depth=3)
+            # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
+            ```
         """
         match_conditions = []
         if prefix:
@@ -645,7 +863,7 @@ __all__ = [
     "SearchOp",
     "ListNamespacesOp",
     "MatchCondition",
-    "NameSpacePath",
+    "NamespacePath",
     "NamespaceMatchType",
     "Embeddings",
     "ensure_embeddings",
diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py
index e6b00efdb..33c502574 100644
--- a/libs/checkpoint/langgraph/store/base/batch.py
+++ b/libs/checkpoint/langgraph/store/base/batch.py
@@ -8,7 +8,7 @@ from langgraph.store.base import (
     Item,
     ListNamespacesOp,
     MatchCondition,
-    NameSpacePath,
+    NamespacePath,
     Op,
     PutOp,
     SearchItem,
@@ -77,8 +77,8 @@ class AsyncBatchedBaseStore(BaseStore):
     async def alist_namespaces(
         self,
         *,
-        prefix: Optional[NameSpacePath] = None,
-        suffix: Optional[NameSpacePath] = None,
+        prefix: Optional[NamespacePath] = None,
+        suffix: Optional[NamespacePath] = None,
         max_depth: Optional[int] = None,
         limit: int = 100,
         offset: int = 0,
diff --git a/libs/checkpoint/langgraph/store/base/embed.py b/libs/checkpoint/langgraph/store/base/embed.py
index 9d869b6ef..8ec33350e 100644
--- a/libs/checkpoint/langgraph/store/base/embed.py
+++ b/libs/checkpoint/langgraph/store/base/embed.py
@@ -28,7 +28,7 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi
 
 
 def ensure_embeddings(
-    embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None],
+    embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc],
 ) -> Embeddings:
     """Ensure that an embedding function conforms to LangChain's Embeddings interface.
 
@@ -44,13 +44,24 @@ def ensure_embeddings(
     Returns:
         An Embeddings instance that wraps the provided function(s).
 
-    Example:
-        >>> def my_embed_fn(texts): return [[0.1, 0.2] for _ in texts]
-        >>> async def my_async_fn(texts): return [[0.1, 0.2] for _ in texts]
-        >>> # Wrap a sync function
-        >>> embeddings = ensure_embeddings(my_embed_fn)
-        >>> # Wrap an async function
-        >>> embeddings = ensure_embeddings(my_async_fn)
+    ??? example "Examples"
+        Wrap a synchronous embedding function:
+        ```python
+        def my_embed_fn(texts):
+            return [[0.1, 0.2] for _ in texts]
+        
+        embeddings = ensure_embeddings(my_embed_fn)
+        result = embeddings.embed_query("hello")  # Returns [0.1, 0.2]
+        ```
+
+        Wrap an asynchronous embedding function:
+        ```python
+        async def my_async_fn(texts):
+            return [[0.1, 0.2] for _ in texts]
+        
+        embeddings = ensure_embeddings(my_async_fn)
+        result = await embeddings.aembed_query("hello")  # Returns [0.1, 0.2]
+        ```
     """
     if embed is None:
         raise ValueError("embed must be provided")
@@ -75,21 +86,27 @@ class EmbeddingsLambda(Embeddings):
             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
+    ??? example "Examples"
+        With a sync function:
+        ```python
+        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:
+        ```python
+        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__(
@@ -179,12 +196,14 @@ def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]:
 
     Args:
         obj: The object to extract text from
-        path: Either a path string or pre-tokenized path list. Path string supports:
-            - Simple paths: "field1.field2"
-            - Array indexing: "[0]", "[*]", "[-1]"
-            - Wildcards: "*"
-            - Multi-field selection: "{field1,field2}"
-            - Nested paths in multi-field: "{field1,nested.field2}"
+        path: Either a path string or pre-tokenized path list. 
+        
+    !!! info "Path types handled"
+        - Simple paths: "field1.field2"
+        - Array indexing: "[0]", "[*]", "[-1]"
+        - Wildcards: "*"
+        - Multi-field selection: "{field1,field2}"
+        - Nested paths in multi-field: "{field1,nested.field2}"
     """
     if not path or path == "$":
         return [json.dumps(obj, sort_keys=True)]
@@ -271,11 +290,11 @@ def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]:
 def tokenize_path(path: str) -> list[str]:
     """Tokenize a path into components.
 
-    Handles:
-    - Simple paths: "field1.field2"
-    - Array indexing: "[0]", "[*]", "[-1]"
-    - Wildcards: "*"
-    - Multi-field selection: "{field1,field2}"
+    !!! info "Types handled"
+        - Simple paths: "field1.field2"
+        - Array indexing: "[0]", "[*]", "[-1]"
+        - Wildcards: "*"
+        - Multi-field selection: "{field1,field2}"
     """
     if not path:
         return []