Rename file

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.
     """
 
 

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

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(

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