Add test

libs/checkpoint/Makefile

@@ -4,11 +4,13 @@
 # TESTING AND COVERAGE
 ######################
 
+TEST ?= .
+
 test:
-	poetry run pytest tests
+	poetry run pytest $(TEST)
 
 test_watch:
-	poetry run ptw .
+	poetry run ptw $(TEST)
 
 ######################
 # LINTING AND FORMATTING

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

@@ -2,8 +2,8 @@
 
 import math
 import random
-from collections import Counter
-from typing import Any, Optional
+from collections import Counter, defaultdict
+from typing import Any
 
 from langchain_core.embeddings import Embeddings
 
@@ -14,36 +14,28 @@ class CharacterEmbeddings(Embeddings):
     def __init__(self, dims: int = 50, seed: int = 42):
         """Initialize with embedding dimensions and random seed."""
         self._rng = random.Random(seed)
-        self._char_to_idx: dict[str, int] = {}
-        self._projection: Optional[list[list[float]]] = None
         self.dims = dims
-
-    def _ensure_projection_matrix(self, texts: list[str]) -> None:
-        """Lazily initialize character mapping and projection matrix."""
-        if self._projection is None:
-            chars = sorted(set("".join(texts)))
-            self._char_to_idx = {c: i for i, c in enumerate(chars)}
-            self._projection = [
-                [self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims)]
-                for _ in range(len(chars))
+        # Create projection vector for each character lazily
+        self._char_projections: defaultdict[str, list[float]] = defaultdict(
+            lambda: [
+                self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims)
             ]
+        )
 
     def _embed_one(self, text: str) -> list[float]:
         """Embed a single text."""
         counts = Counter(text)
-        char_vec = [0.0] * len(self._char_to_idx)
+        total = sum(counts.values())
 
+        if total == 0:
+            return [0.0] * self.dims
+
+        embedding = [0.0] * self.dims
         for char, count in counts.items():
-            if char in self._char_to_idx:
-                char_vec[self._char_to_idx[char]] = count
-
-        total = sum(char_vec)
-        if total > 0:
-            char_vec = [v / total for v in char_vec]
-        embedding = [
-            sum(a * b for a, b in zip(char_vec, proj))
-            for proj in zip(*self._projection)
-        ]
+            weight = count / total
+            char_proj = self._char_projections[char]
+            for i, proj in enumerate(char_proj):
+                embedding[i] += weight * proj
 
         norm = math.sqrt(sum(x * x for x in embedding))
         if norm > 0:
@@ -53,12 +45,10 @@ class CharacterEmbeddings(Embeddings):
 
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
         """Embed a list of documents."""
-        self._ensure_projection_matrix(texts)
         return [self._embed_one(text) for text in texts]
 
     def embed_query(self, text: str) -> list[float]:
         """Embed a query string."""
-        self._ensure_projection_matrix([text])
         return self._embed_one(text)
 
     def __eq__(self, other: Any) -> bool:

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

@@ -199,7 +199,7 @@ class InMemoryStore(BaseStore):
                 with cf.ThreadPoolExecutor() as executor:
                     futures = {
                         q: executor.submit(self.embeddings.embed_query, q)
-                        for q in queries
+                        for q in list(queries)
                     }
                     for query, future in futures.items():
                         queryinmem_store[query] = future.result()
@@ -215,7 +215,7 @@ class InMemoryStore(BaseStore):
             queries = {op.query for (op, _) in search_ops.values() if op.query}
 
             if queries:
-                coros = [self.embeddings.aembed_query(q) for q in queries]
+                coros = [self.embeddings.aembed_query(q) for q in list(queries)]
                 results = await asyncio.gather(*coros)
                 queryinmem_store = dict(zip(queries, results))
 

libs/checkpoint/tests/test_store.py

@@ -881,3 +881,53 @@ async def test_async_vector_search_edge_cases(
     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
+    store = InMemoryStore(
+        embedding_config={
+            "dims": fake_embeddings.dims,
+            "embed": fake_embeddings,
+            # Key 2 isn't included. Don't index it.
+            "text_fields": ["key0", "key1", "key3"],
+        }
+    )
+    # This will have 2 vectors representing it
+    doc1 = {
+        # Omit key0 - check it doesn't raise an error
+        "key1": "xxx",
+        "key2": "yyy",
+        "key3": "zzz",
+    }
+    # This will have 3 vectors representing it
+    doc2 = {
+        "key0": "uuu",
+        "key1": "vvv",
+        "key2": "www",
+        "key3": "xxx",
+    }
+    await store.aput(("test",), "doc1", doc1)
+    await store.aput(("test",), "doc2", doc2)
+
+    # doc2.key3 and doc1.key1 both would have the highest score
+    results = await store.asearch(("test",), query="xxx")
+    assert len(results) == 2
+    assert results[0].key != results[1].key
+    ascore = results[0].response_metadata["score"]
+    bscore = results[1].response_metadata["score"]
+    assert ascore == bscore
+
+    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)
+
+    # 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