From 9a83dd79003552931ce1f600801ce047fee09b0f Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 26 Nov 2024 19:23:48 -0800 Subject: [PATCH] Add test --- libs/checkpoint/Makefile | 6 ++- .../langgraph/store/base/_embed_test_utils.py | 42 ++++++---------- .../langgraph/store/memory/__init__.py | 4 +- libs/checkpoint/tests/test_store.py | 50 +++++++++++++++++++ 4 files changed, 72 insertions(+), 30 deletions(-) diff --git a/libs/checkpoint/Makefile b/libs/checkpoint/Makefile index ddf087ef5..2636db8fe 100644 --- a/libs/checkpoint/Makefile +++ b/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 diff --git a/libs/checkpoint/langgraph/store/base/_embed_test_utils.py b/libs/checkpoint/langgraph/store/base/_embed_test_utils.py index 10e1e373c..d28cd959f 100644 --- a/libs/checkpoint/langgraph/store/base/_embed_test_utils.py +++ b/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: diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py index cfcf03fff..b350c8a17 100644 --- a/libs/checkpoint/langgraph/store/memory/__init__.py +++ b/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)) diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index fee7d1ccb..8c5df5b04 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/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