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