diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 4f0cb8daa..22a28b8c2 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -37,7 +37,6 @@ from langgraph.store.base import ( ListNamespacesOp, Op, PutOp, - ResponseMetadata, Result, SearchItem, SearchOp, @@ -903,20 +902,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 ef6729d83..033cfba12 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,97 +119,350 @@ 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 class GetOp(NamedTuple): - """Operation to retrieve an item by namespace and key.""" + """Operation to retrieve a specific item by its namespace and key. + + This operation allows precise retrieval of stored items using their full path + (namespace) and unique identifier (key) combination. + + ??? example "Examples" + + Basic item retrieval: + ```python + GetOp(namespace=("users", "profiles"), key="user123") + GetOp(namespace=("cache", "embeddings"), key="doc456") + ``` + """ namespace: tuple[str, ...] - """Hierarchical path for the item.""" + """Hierarchical path that uniquely identifies the item's location. + + ??? example "Examples" + + ```python + ("users",) # Root level users namespace + ("users", "profiles") # Profiles within users namespace + ``` + """ + key: str - """Unique identifier within the namespace.""" + """Unique identifier for the item within its specific namespace. + + ??? example "Examples" + + ```python + "user123" # For a user profile + "doc456" # For a document + ``` + """ class SearchOp(NamedTuple): - """Operation to search for items within a namespace prefix.""" + """Operation to search for items within a specified namespace hierarchy. + + This operation supports both structured filtering and natural language search + within a given namespace prefix. It provides pagination through limit and offset + parameters. + + Note: + Natural language search support depends on your store implementation. + + ??? example "Examples" + Search with filters and pagination: + ```python + SearchOp( + namespace_prefix=("documents",), + filter={"type": "report", "status": "active"}, + limit=5, + offset=10 + ) + ``` + + Natural language search: + ```python + SearchOp( + namespace_prefix=("users", "content"), + query="technical documentation about APIs", + limit=20 + ) + ``` + """ namespace_prefix: tuple[str, ...] - """Hierarchical path prefix to search within.""" + """Hierarchical path prefix defining the search scope. + + ??? example "Examples" + + ```python + () # Search entire store + ("documents",) # Search all documents + ("users", "content") # Search within user content + ``` + """ + filter: Optional[dict[str, Any]] = None - """Key-value pairs to filter results.""" + """Key-value pairs for filtering results based on exact matches or comparison operators. + + The filter supports both exact matches and operator-based comparisons. + + Supported Operators: + - $eq: Equal to (same as direct value comparison) + - $ne: Not equal to + - $gt: Greater than + - $gte: Greater than or equal to + - $lt: Less than + - $lte: Less than or equal to + + ??? example "Examples" + + Simple exact match: + + ```python + {"status": "active"} + ``` + + Comparison operators: + + ```python + {"score": {"$gt": 4.99}} # Score greater than 4.99 + ``` + + Multiple conditions: + + ```python + { + "score": {"$gte": 3.0}, + "color": "red" + } + ``` + + Note: + Comparison operator support depends on your store implementation. + """ + limit: int = 10 - """Maximum number of items to return.""" + """Maximum number of items to return in the search results.""" + offset: int = 0 - """Number of items to skip before returning results.""" + """Number of matching items to skip for pagination.""" + query: Optional[str] = None - """The search query for natural language search.""" + """Natural language search query for semantic search capabilities. - -class PutOp(NamedTuple): - """Operation to store, update, or delete an item.""" - - namespace: tuple[str, ...] - """Hierarchical path for the item. - - Represented as a tuple of strings, allowing for nested categorization. - For example: ("documents", "user123") - """ - - key: str - """Unique identifier for the document. - - Should be distinct within its namespace. - """ - - value: Optional[dict[str, Any]] - """Data to be stored, or None to delete the item. - - 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 - """ - 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. + ??? example "Examples" + - "technical documentation about REST APIs" + - "machine learning papers from 2023" """ -NameSpacePath = tuple[Union[str, Literal["*"]], ...] +# Type representing a namespace path that can include wildcards +NamespacePath = tuple[Union[str, Literal["*"]], ...] +"""A tuple representing a namespace path that can include wildcards. +Examples: + ("users",) # Exact users namespace + ("documents", "*") # Any sub-namespace under documents + ("cache", "*", "v1") # Any cache category with v1 version +""" + +# Type for specifying how to match namespaces NamespaceMatchType = Literal["prefix", "suffix"] +"""Specifies how to match namespace paths. + +Values: + "prefix": Match from the start of the namespace + "suffix": Match from the end of the namespace +""" class MatchCondition(NamedTuple): - """Represents a single match condition.""" + """Represents a pattern for matching namespaces in the store. + + This class combines a match type (prefix or suffix) with a namespace path + pattern that can include wildcards to flexibly match different namespace + hierarchies. + + ??? example "Examples" + Prefix matching: + ```python + MatchCondition(match_type="prefix", path=("users", "profiles")) + ``` + + Suffix matching with wildcard: + ```python + MatchCondition(match_type="suffix", path=("cache", "*")) + ``` + + Simple suffix matching: + ```python + MatchCondition(match_type="suffix", path=("v1",)) + ``` + """ match_type: NamespaceMatchType - path: NameSpacePath + """Type of namespace matching to perform.""" + + path: NamespacePath + """Namespace path pattern that can include wildcards.""" class ListNamespacesOp(NamedTuple): - """Operation to list namespaces with optional match conditions.""" + """Operation to list and filter namespaces in the store. + + This operation allows exploring the organization of data, finding specific + collections, and navigating the namespace hierarchy. + + ??? example "Examples" + + List all namespaces under the "documents" path: + ```python + ListNamespacesOp( + match_conditions=(MatchCondition(match_type="prefix", path=("documents",)),), + max_depth=2 + ) + ``` + + List all namespaces that end with "v1": + ```python + ListNamespacesOp( + match_conditions=(MatchCondition(match_type="suffix", path=("v1",)),), + limit=50 + ) + ``` + + """ match_conditions: Optional[tuple[MatchCondition, ...]] = None - """A tuple of match conditions to apply to namespaces.""" + """Optional conditions for filtering namespaces. + + ??? example "Examples" + All user namespaces: + ```python + (MatchCondition(match_type="prefix", path=("users",)),) + ``` + + All namespaces that start with "docs" and end with "draft": + ```python + ( + MatchCondition(match_type="prefix", path=("docs",)), + MatchCondition(match_type="suffix", path=("draft",)) + ) + ``` + """ max_depth: Optional[int] = None - """Return namespaces up to this depth in the hierarchy.""" + """Maximum depth of namespace hierarchy to return. + + Note: + Namespaces deeper than this level will be truncated. + """ limit: int = 100 """Maximum number of namespaces to return.""" offset: int = 0 - """Number of namespaces to skip before returning results.""" + """Number of namespaces to skip for pagination.""" + + +class PutOp(NamedTuple): + """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 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. + + ??? example "Examples" + Root level documents + ```python + ("documents",) + ``` + + User-specific documents + ```python + ("documents", "user123") + ``` + + Nested cache structure + ```python + ("cache", "embeddings", "v1") + ``` + """ + + key: str + """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]] + """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. + + Example: + { + "field1": "string value", + "field2": 123, + "nested": {"can": "contain", "any": "serializable data"} + } + """ + + 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[*]" + + ??? example "Examples" + - None - Use store defaults + - False - Don't index this item + - list[str] - List of fields to index + + ```python + [ + "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 + ] + ``` + """ Op = Union[GetOp, SearchOp, PutOp, ListNamespacesOp] @@ -227,8 +473,8 @@ class InvalidNamespaceError(ValueError): """Provided namespace is invalid.""" -class EmbeddingConfig(TypedDict, total=False): - """Configuration for vector embeddings in PostgreSQL store.""" +class IndexConfig(TypedDict, total=False): + """Configuration for indexing documents for semantic search in the store.""" dims: int """Number of dimensions in the embedding vectors. @@ -245,17 +491,11 @@ class EmbeddingConfig(TypedDict, total=False): embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc] """Optional function to generate embeddings from text.""" - aembed: Optional[AEmbeddingsFunc] - """Optional asynchronous function to generate embeddings from text. - - Provide for asynchronous embedding generation if you do not provide - an Embeddings object. - """ - 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. """ @@ -333,16 +573,44 @@ 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. + + ??? example "Examples" + Simple storage without special indexing (respects store defaults) + ```python + store.put(("docs",), "report", {"title": "Annual Report"}) + ``` + + Index specific fields for search + ```python + store.put(("docs",), "report", {"title": "Annual Report"}, index=["title"]) + ``` + + Do not index for semantic search + ```python + store.put(("docs",), "report", {"title": "Annual Report"}, index=False) + ``` """ _validate_namespace(namespace) self.batch([PutOp(namespace, key, value, index=index)]) @@ -359,8 +627,8 @@ class BaseStore(ABC): def list_namespaces( self, *, - prefix: Optional[NameSpacePath] = None, - suffix: Optional[NameSpacePath] = None, + prefix: Optional[NamespacePath] = None, + suffix: Optional[NamespacePath] = None, max_depth: Optional[int] = None, limit: int = 100, offset: int = 0, @@ -374,7 +642,7 @@ class BaseStore(ABC): prefix (Optional[Tuple[str, ...]]): Filter namespaces that start with this path. suffix (Optional[Tuple[str, ...]]): Filter namespaces that end with this path. max_depth (Optional[int]): Return namespaces up to this depth in the hierarchy. - Namespaces deeper than this level will be truncated to this depth. + Namespaces deeper than this level will be truncated. limit (int): Maximum number of namespaces to return (default 100). offset (int): Number of namespaces to skip for pagination (default 0). @@ -382,16 +650,18 @@ class BaseStore(ABC): List[Tuple[str, ...]]: A list of namespace tuples that match the criteria. Each tuple represents a full namespace path up to `max_depth`. - Examples: - + ??? example "Examples": Setting max_depth=3. Given the namespaces: - # ("a", "b", "c") - # ("a", "b", "d", "e") - # ("a", "b", "d", "i") - # ("a", "b", "f") - # ("a", "c", "f") - store.list_namespaces(prefix=("a", "b"), max_depth=3) - # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] + ```python + # Example if you have the following namespaces: + # ("a", "b", "c") + # ("a", "b", "d", "e") + # ("a", "b", "d", "i") + # ("a", "b", "f") + # ("a", "c", "f") + store.list_namespaces(prefix=("a", "b"), max_depth=3) + # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] + ``` """ match_conditions = [] if prefix: @@ -452,19 +722,50 @@ 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. + + ??? example "Examples" + Simple storage without special indexing: + ```python + await store.aput(("docs",), "report", {"title": "Annual Report"}) + ``` + + Index specific fields for search: + ```python + 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. @@ -478,8 +779,8 @@ class BaseStore(ABC): async def alist_namespaces( self, *, - prefix: Optional[NameSpacePath] = None, - suffix: Optional[NameSpacePath] = None, + prefix: Optional[NamespacePath] = None, + suffix: Optional[NamespacePath] = None, max_depth: Optional[int] = None, limit: int = 100, offset: int = 0, @@ -501,16 +802,19 @@ class BaseStore(ABC): List[Tuple[str, ...]]: A list of namespace tuples that match the criteria. Each tuple represents a full namespace path up to `max_depth`. - Examples: + ??? example "Examples" + Setting max_depth=3 with existing namespaces: + ```python + # Given the following namespaces: + # ("a", "b", "c") + # ("a", "b", "d", "e") + # ("a", "b", "d", "i") + # ("a", "b", "f") + # ("a", "c", "f") - Setting max_depth=3. Given the namespaces: - # ("a", "b", "c") - # ("a", "b", "d", "e") - # ("a", "b", "d", "i") - # ("a", "b", "f") - # ("a", "c", "f") - await store.alist_namespaces(prefix=("a", "b"), max_depth=3) - # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] + await store.alist_namespaces(prefix=("a", "b"), max_depth=3) + # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] + ``` """ match_conditions = [] if prefix: @@ -559,7 +863,7 @@ __all__ = [ "SearchOp", "ListNamespacesOp", "MatchCondition", - "NameSpacePath", + "NamespacePath", "NamespaceMatchType", "Embeddings", "ensure_embeddings", diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index bb74eae8e..33c502574 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, @@ -8,7 +8,7 @@ from langgraph.store.base import ( Item, ListNamespacesOp, MatchCondition, - NameSpacePath, + NamespacePath, Op, PutOp, SearchItem, @@ -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() @@ -77,8 +77,8 @@ class AsyncBatchedBaseStore(BaseStore): async def alist_namespaces( self, *, - prefix: Optional[NameSpacePath] = None, - suffix: Optional[NameSpacePath] = None, + prefix: Optional[NamespacePath] = None, + suffix: Optional[NamespacePath] = None, max_depth: Optional[int] = None, limit: int = 100, offset: int = 0, diff --git a/libs/checkpoint/langgraph/store/base/_embed.py b/libs/checkpoint/langgraph/store/base/embed.py similarity index 79% rename from libs/checkpoint/langgraph/store/base/_embed.py rename to libs/checkpoint/langgraph/store/base/embed.py index a04b4f138..8ec33350e 100644 --- a/libs/checkpoint/langgraph/store/base/_embed.py +++ b/libs/checkpoint/langgraph/store/base/embed.py @@ -28,9 +28,7 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi def ensure_embeddings( - embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None], - *, - aembed: Optional[AEmbeddingsFunc] = None, + embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc], ) -> Embeddings: """Ensure that an embedding function conforms to LangChain's Embeddings interface. @@ -42,73 +40,85 @@ def ensure_embeddings( embed: Either an existing Embeddings instance, or a function that converts text to embeddings. If the function is async, it will be used for both sync and async operations. - aembed: Optional async function for embeddings. If provided, it will be used - for async operations while the sync function is used for sync operations. - Must be None if embed is async. Returns: An Embeddings instance that wraps the provided function(s). - Example: - >>> def my_embed_fn(texts): return [[0.1, 0.2] for _ in texts] - >>> async def my_async_fn(texts): return [[0.1, 0.2] for _ in texts] - >>> # Wrap a sync function - >>> embeddings = ensure_embeddings(my_embed_fn) - >>> # Wrap an async function - >>> embeddings = ensure_embeddings(my_async_fn) - >>> # Provide both sync and async implementations - >>> embeddings = ensure_embeddings(my_embed_fn, aembed=my_async_fn) + ??? example "Examples" + Wrap a synchronous embedding function: + ```python + def my_embed_fn(texts): + return [[0.1, 0.2] for _ in texts] + + embeddings = ensure_embeddings(my_embed_fn) + result = embeddings.embed_query("hello") # Returns [0.1, 0.2] + ``` + + Wrap an asynchronous embedding function: + ```python + async def my_async_fn(texts): + return [[0.1, 0.2] for _ in texts] + + embeddings = ensure_embeddings(my_async_fn) + result = await embeddings.aembed_query("hello") # Returns [0.1, 0.2] + ``` """ - if embed is None and aembed is None: - raise ValueError("embed or aembed must be provided") + if embed is None: + raise ValueError("embed must be provided") if isinstance(embed, Embeddings): return embed - return EmbeddingsLambda(embed, afunc=aembed) + return EmbeddingsLambda(embed) 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: - >>> 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] + ??? example "Examples" + With a sync function: + ```python + 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: + ```python + 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], - afunc: Optional[AEmbeddingsFunc] = None, + func: Union[EmbeddingsFunc, AEmbeddingsFunc], ) -> None: + if func is None: + raise ValueError("func must be provided") if _is_async_callable(func): - if afunc is not None: - raise ValueError( - "afunc must be None if func is async. The async func will be used for both sync and async operations." - ) self.afunc = func else: self.func = func - if afunc is not None: - self.afunc = afunc def embed_documents(self, texts: list[str]) -> list[list[float]]: """Embed a list of texts into vectors. @@ -186,14 +196,16 @@ def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]: Args: obj: The object to extract text from - path: Either a path string or pre-tokenized path list. Path string supports: - - Simple paths: "field1.field2" - - Array indexing: "[0]", "[*]", "[-1]" - - Wildcards: "*" - - Multi-field selection: "{field1,field2}" - - Nested paths in multi-field: "{field1,nested.field2}" + path: Either a path string or pre-tokenized path list. + + !!! info "Path types handled" + - Simple paths: "field1.field2" + - Array indexing: "[0]", "[*]", "[-1]" + - Wildcards: "*" + - 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 @@ -278,11 +290,11 @@ def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]: def tokenize_path(path: str) -> list[str]: """Tokenize a path into components. - Handles: - - Simple paths: "field1.field2" - - Array indexing: "[0]", "[*]", "[-1]" - - Wildcards: "*" - - Multi-field selection: "{field1,field2}" + !!! info "Types handled" + - Simple paths: "field1.field2" + - Array indexing: "[0]", "[*]", "[-1]" + - Wildcards: "*" + - Multi-field selection: "{field1,field2}" """ if not path: return [] diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py index b350c8a17..40011a7ba 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, - EmbeddingConfig, 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[EmbeddingConfig] = 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: @@ -232,20 +230,24 @@ class InMemoryStore(BaseStore): if not candidates: results[i] = [] continue - if op.query: + 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: @@ -258,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( @@ -266,7 +274,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) if score is not None else None, ) for score, item in kept ] @@ -315,7 +323,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 +336,16 @@ 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"]: + 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: @@ -361,7 +373,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 8c5df5b04..053d9a981 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) ] @@ -382,12 +383,14 @@ async def test_cannot_put_empty_namespace() -> None: await store.aput(("foo", "langgraph", "foo"), "bar", doc) assert (await store.aget(("foo", "langgraph", "foo"), "bar")).value == doc # type: ignore[union-attr] - assert (await store.asearch(("foo", "langgraph", "foo")))[0].value == doc + assert (await store.asearch(("foo", "langgraph", "foo"), query="bar"))[ + 0 + ].value == doc await store.adelete(("foo", "langgraph", "foo"), "bar") 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 @@ -423,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 @@ -508,11 +514,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( @@ -520,7 +526,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"}), @@ -547,7 +553,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"}), @@ -572,7 +578,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"}) @@ -581,20 +587,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) @@ -607,7 +613,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"}) @@ -616,20 +622,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) @@ -640,7 +646,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 = [ @@ -680,7 +686,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 = [ @@ -720,7 +726,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"] @@ -800,7 +806,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}"}) @@ -821,7 +827,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}"}) @@ -837,60 +843,14 @@ 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( - embedding_config={"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( - embedding_config={"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( - 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 @@ -914,20 +874,76 @@ 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 + + # 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