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] 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