diff --git a/libs/checkpoint-duckdb/langgraph/store/duckdb/base.py b/libs/checkpoint-duckdb/langgraph/store/duckdb/base.py index e0fb57067..89bf13681 100644 --- a/libs/checkpoint-duckdb/langgraph/store/duckdb/base.py +++ b/libs/checkpoint-duckdb/langgraph/store/duckdb/base.py @@ -23,6 +23,7 @@ from langgraph.store.base import ( Op, PutOp, Result, + SearchItem, SearchOp, ) @@ -283,7 +284,7 @@ class DuckDBStore(BaseStore, BaseDuckDBStore[duckdb.DuckDBPyConnection]): for cur, idx in cursors: rows = cur.fetchall() - items = [_row_to_item(_convert_ns(row[0]), row) for row in rows] + items = [_row_to_search_item(_convert_ns(row[0]), row) for row in rows] results[idx] = items def _batch_list_namespaces_ops( @@ -376,6 +377,22 @@ def _row_to_item( ) +def _row_to_search_item( + namespace: tuple[str, ...], + row: tuple, +) -> SearchItem: + """Convert a row from the database into an SearchItem.""" + # TODO: Add support for search + _, key, val, created_at, updated_at = row + return SearchItem( + value=val if isinstance(val, dict) else json.loads(val), + key=key, + namespace=namespace, + created_at=created_at, + updated_at=updated_at, + ) + + def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]: grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list) tot = 0 diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 2107b05a3..1a3aff119 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -378,15 +378,19 @@ class PostgresSaver(BasePostgresSaver): # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock if self.supports_pipeline: - with self.lock, conn.pipeline(), conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + with ( + self.lock, + conn.pipeline(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur else: # Use connection's transaction context manager when pipeline mode not supported - with self.lock, conn.transaction(), conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + with ( + self.lock, + conn.transaction(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur else: with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur: diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 589520efc..5b07e0067 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -338,20 +338,25 @@ class AsyncPostgresSaver(BasePostgresSaver): # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock if self.supports_pipeline: - async with self.lock, conn.pipeline(), conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + async with ( + self.lock, + conn.pipeline(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur else: # Use connection's transaction context manager when pipeline mode not supported - async with self.lock, conn.transaction(), conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + async with ( + self.lock, + conn.transaction(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur else: - async with self.lock, conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + async with ( + self.lock, + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur def list( diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index 578523052..b90a9a0d5 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -28,6 +28,7 @@ from langgraph.store.postgres.base import ( _decode_ns_bytes, _group_ops, _row_to_item, + _row_to_search_item, ) logger = logging.getLogger(__name__) @@ -146,7 +147,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con await cur.execute(query, params) rows = cast(list[Row], await cur.fetchall()) items = [ - _row_to_item( + _row_to_search_item( _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer ) for row in rows @@ -195,9 +196,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con async with self.lock, conn.pipeline(), conn.cursor(binary=True) as cur: yield cur else: - async with self.lock, conn.transaction(), conn.cursor( - binary=True - ) as cur: + async with ( + self.lock, + conn.transaction(), + conn.cursor(binary=True) as cur, + ): yield cur else: async with conn.cursor(binary=True) as cur: diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 3bb343b0c..200218c0d 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -36,6 +36,7 @@ from langgraph.store.base import ( Op, PutOp, Result, + SearchItem, SearchOp, ) @@ -344,14 +345,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock if self.supports_pipeline: - with self.lock, conn.pipeline(), conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + with ( + self.lock, + conn.pipeline(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur else: - with self.lock, conn.transaction(), conn.cursor( - binary=True, row_factory=dict_row - ) as cur: + with ( + self.lock, + conn.transaction(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): yield cur else: with conn.cursor(binary=True, row_factory=dict_row) as cur: @@ -430,7 +435,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): cur.execute(query, params) rows = cast(list[Row], cur.fetchall()) results[idx] = [ - _row_to_item( + _row_to_search_item( _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer ) for row in rows @@ -517,6 +522,32 @@ def _row_to_item( ) +def _row_to_search_item( + namespace: tuple[str, ...], + row: Row, + *, + loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None, +) -> SearchItem: + """Convert a row from the database into an Item.""" + loader = loader or _json_loads + val = row["value"] + 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"], + score=score, + ) + + def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]: grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list) tot = 0 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/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index 098462339..033cfba12 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -1,12 +1,27 @@ """Base classes and types for persistent key-value stores. -Stores enable persistence and memory that can be shared across threads, -scoped to user IDs, assistant IDs, or other arbitrary namespaces. +Stores provide long-term memory that persists across threads and conversations. +Supports hierarchical namespaces, key-value storage, and optional vector search. + +Core types: +- BaseStore: Store interface with sync/async operations +- Item: Stored key-value pairs with metadata +- Op: Get/Put/Search/List operations """ from abc import ABC, abstractmethod from datetime import datetime -from typing import Any, Iterable, Literal, NamedTuple, Optional, Union, cast +from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast + +from langchain_core.embeddings import Embeddings + +from langgraph.store.base.embed import ( + AEmbeddingsFunc, + EmbeddingsFunc, + ensure_embeddings, + get_text_at_path, + tokenize_path, +) class Item: @@ -73,112 +88,415 @@ class Item: } +class SearchItem(Item): + """Represents a result item with additional response metadata.""" + + __slots__ = ("score",) + + def __init__( + self, + namespace: tuple[str, ...], + key: str, + value: dict[str, Any], + created_at: datetime, + updated_at: datetime, + score: Optional[float] = None, + ) -> None: + """Initialize a result item. + + Args: + namespace: Hierarchical path to the item. + key: Unique identifier within the namespace. + value: The stored value. + created_at: When the item was first created. + updated_at: When the item was last updated. + score: Relevance/similarity score if from a ranked operation. + """ + super().__init__( + value=value, + key=key, + namespace=namespace, + created_at=created_at, + updated_at=updated_at, + ) + self.score = score + + def dict(self) -> dict: + result = super().dict() + 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 + """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 + ??? 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] -Result = Union[Item, list[Item], list[tuple[str, ...]], None] +Result = Union[Item, list[Item], list[SearchItem], list[tuple[str, ...]], None] class InvalidNamespaceError(ValueError): """Provided namespace is invalid.""" -def _validate_namespace(namespace: tuple[str, ...]) -> None: - if not namespace: - raise InvalidNamespaceError("Namespace cannot be empty.") - for label in namespace: - if not isinstance(label, str): - raise InvalidNamespaceError( - f"Invalid namespace label '{label}' found in {namespace}. Namespace labels" - f" must be strings, but got {type(label).__name__}." - ) - if "." in label: - raise InvalidNamespaceError( - f"Invalid namespace label '{label}' found in {namespace}. Namespace labels cannot contain periods ('.')." - ) - elif not label: - raise InvalidNamespaceError( - f"Namespace labels cannot be empty strings. Got {label} in {namespace}" - ) - if namespace[0] == "langgraph": - raise InvalidNamespaceError( - f'Root label for namespace cannot be "langgraph". Got: {namespace}' - ) +class IndexConfig(TypedDict, total=False): + """Configuration for indexing documents for semantic search in the store.""" + + dims: int + """Number of dimensions in the embedding vectors. + + Common embedding models have the following dimensions: + - OpenAI text-embedding-3-large: 256, 1024, or 3072 + - OpenAI text-embedding-3-small: 512 or 1536 + - OpenAI text-embedding-ada-002: 1536 + - Cohere embed-english-v3.0: 1024 + - Cohere embed-english-light-v3.0: 384 + - Cohere embed-multilingual-v3.0: 1024 + - Cohere embed-multilingual-light-v3.0: 384 + """ + + embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc] + """Optional function to generate embeddings from text.""" + + fields: Optional[list[str]] + """Fields to extract text from for embedding generation. + + Defaults to the root ["$"], which embeds the json object as a whole. + """ class BaseStore(ABC): @@ -231,14 +549,16 @@ class BaseStore(ABC): namespace_prefix: tuple[str, ...], /, *, + query: Optional[str] = None, filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - ) -> list[Item]: + ) -> list[SearchItem]: """Search for items within a namespace prefix. Args: namespace_prefix: Hierarchical path prefix to search within. + query: Optional query for natural language search. filter: Key-value pairs to filter results. limit: Maximum number of items to return. offset: Number of items to skip before returning results. @@ -246,18 +566,54 @@ class BaseStore(ABC): Returns: List of items matching the search criteria. """ - return self.batch([SearchOp(namespace_prefix, filter, limit, offset)])[0] + return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0] - def put(self, namespace: tuple[str, ...], key: str, value: dict[str, Any]) -> None: - """Store or update an item. + def put( + self, + namespace: tuple[str, ...], + key: str, + value: dict[str, Any], + index: Optional[Union[Literal[False], list[str]]] = None, + ) -> None: + """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. + 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)]) + self.batch([PutOp(namespace, key, value, index=index)]) def delete(self, namespace: tuple[str, ...], key: str) -> None: """Delete an item. @@ -271,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, @@ -286,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). @@ -294,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: @@ -336,14 +694,16 @@ class BaseStore(ABC): namespace_prefix: tuple[str, ...], /, *, + query: Optional[str] = None, filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - ) -> list[Item]: + ) -> list[SearchItem]: """Asynchronously search for items within a namespace prefix. Args: namespace_prefix: Hierarchical path prefix to search within. + query: Optional query for natural language search. filter: Key-value pairs to filter results. limit: Maximum number of items to return. offset: Number of items to skip before returning results. @@ -351,22 +711,61 @@ class BaseStore(ABC): Returns: List of items matching the search criteria. """ - return (await self.abatch([SearchOp(namespace_prefix, filter, limit, offset)]))[ - 0 - ] + return ( + await self.abatch( + [SearchOp(namespace_prefix, filter, limit, offset, query)] + ) + )[0] async def aput( - self, namespace: tuple[str, ...], key: str, value: dict[str, Any] + self, + namespace: tuple[str, ...], + key: str, + value: dict[str, Any], + 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. + 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)]) + await self.abatch([PutOp(namespace, key, value, index=index)]) async def adelete(self, namespace: tuple[str, ...], key: str) -> None: """Asynchronously delete an item. @@ -380,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, @@ -403,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: @@ -427,3 +829,44 @@ class BaseStore(ABC): offset=offset, ) return (await self.abatch([op]))[0] + + +def _validate_namespace(namespace: tuple[str, ...]) -> None: + if not namespace: + raise InvalidNamespaceError("Namespace cannot be empty.") + for label in namespace: + if not isinstance(label, str): + raise InvalidNamespaceError( + f"Invalid namespace label '{label}' found in {namespace}. Namespace labels" + f" must be strings, but got {type(label).__name__}." + ) + if "." in label: + raise InvalidNamespaceError( + f"Invalid namespace label '{label}' found in {namespace}. Namespace labels cannot contain periods ('.')." + ) + elif not label: + raise InvalidNamespaceError( + f"Namespace labels cannot be empty strings. Got {label} in {namespace}" + ) + if namespace[0] == "langgraph": + raise InvalidNamespaceError( + f'Root label for namespace cannot be "langgraph". Got: {namespace}' + ) + + +__all__ = [ + "BaseStore", + "Item", + "Op", + "PutOp", + "GetOp", + "SearchOp", + "ListNamespacesOp", + "MatchCondition", + "NamespacePath", + "NamespaceMatchType", + "Embeddings", + "ensure_embeddings", + "tokenize_path", + "get_text_at_path", +] diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index b1030942d..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,9 +8,10 @@ from langgraph.store.base import ( Item, ListNamespacesOp, MatchCondition, - NameSpacePath, + NamespacePath, Op, PutOp, + SearchItem, SearchOp, _validate_namespace, ) @@ -43,12 +44,13 @@ class AsyncBatchedBaseStore(BaseStore): namespace_prefix: tuple[str, ...], /, *, + query: Optional[str] = None, filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - ) -> list[Item]: + ) -> list[SearchItem]: fut = self._loop.create_future() - self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset) + self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query) return await fut async def aput( @@ -56,10 +58,11 @@ class AsyncBatchedBaseStore(BaseStore): namespace: tuple[str, ...], key: str, value: dict[str, Any], + index: Optional[Union[Literal[False], list[str]]] = None, ) -> None: _validate_namespace(namespace) fut = self._loop.create_future() - self._aqueue[fut] = PutOp(namespace, key, value) + self._aqueue[fut] = PutOp(namespace, key, value, index) return await fut async def adelete( @@ -74,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 new file mode 100644 index 000000000..0434481cd --- /dev/null +++ b/libs/checkpoint/langgraph/store/base/embed.py @@ -0,0 +1,380 @@ +"""Utilities for working with embedding functions and LangChain's Embeddings interface. + +This module provides tools to wrap arbitrary embedding functions (both sync and async) +into LangChain's Embeddings interface. This enables using custom embedding functions +with LangChain-compatible tools while maintaining support for both synchronous and +asynchronous operations. +""" + +import asyncio +import json +from typing import Any, Awaitable, Callable, Optional, Sequence, Union + +from langchain_core.embeddings import Embeddings + +EmbeddingsFunc = Callable[[Sequence[str]], list[list[float]]] +"""Type for synchronous embedding functions. + +The function should take a sequence of strings and return a list of embeddings, +where each embedding is a list of floats. The dimensionality of the embeddings +should be consistent for all inputs. +""" + +AEmbeddingsFunc = Callable[[Sequence[str]], Awaitable[list[list[float]]]] +"""Type for asynchronous embedding functions. + +Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddings. +""" + + +def ensure_embeddings( + embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None], +) -> Embeddings: + """Ensure that an embedding function conforms to LangChain's Embeddings interface. + + This function wraps arbitrary embedding functions to make them compatible with + LangChain's Embeddings interface. It handles both synchronous and asynchronous + functions. + + Args: + 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. + + Returns: + An Embeddings instance that wraps the provided function(s). + + ??? 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: + raise ValueError("embed must be provided") + if isinstance(embed, Embeddings): + return embed + 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 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 async operations, but sync operations + will raise an error. If sync, it will be used for both sync and async operations. + + ??? 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: + if func is None: + raise ValueError("func must be provided") + if _is_async_callable(func): + self.afunc = func + else: + self.func = func + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + """Embed a list of texts into vectors. + + Args: + texts: list of texts to convert to embeddings. + + Returns: + list of embeddings, one per input text. Each embedding is a list of floats. + + Raises: + ValueError: If the instance was initialized with only an async function. + """ + func = getattr(self, "func", None) + if func is None: + raise ValueError( + "EmbeddingsLambda was initialized with an async function but no sync function. " + "Use aembed_documents for async operation or provide a sync function." + ) + return func(texts) + + def embed_query(self, text: str) -> list[float]: + """Embed a single piece of text. + + Args: + text: Text to convert to an embedding. + + Returns: + Embedding vector as a list of floats. + + Note: + This is equivalent to calling embed_documents with a single text + and taking the first result. + """ + return self.embed_documents([text])[0] + + async def aembed_documents(self, texts: list[str]) -> list[list[float]]: + """Asynchronously embed a list of texts into vectors. + + Args: + texts: list of texts to convert to embeddings. + + Returns: + list of embeddings, one per input text. Each embedding is a list of floats. + + Note: + If no async function was provided, this falls back to the sync implementation. + """ + afunc = getattr(self, "afunc", None) + if afunc is None: + return await super().aembed_documents(texts) + return await afunc(texts) + + async def aembed_query(self, text: str) -> list[float]: + """Asynchronously embed a single piece of text. + + Args: + text: Text to convert to an embedding. + + Returns: + Embedding vector as a list of floats. + + Note: + This is equivalent to calling aembed_documents with a single text + and taking the first result. + """ + afunc = getattr(self, "afunc", None) + if afunc is None: + return await super().aembed_query(text) + return (await afunc([text]))[0] + + +def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]: + """Extract text from an object using a path expression or pre-tokenized path. + + Args: + obj: The object to extract text from + 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 == "$": + return [json.dumps(obj, sort_keys=True)] + + tokens = tokenize_path(path) if isinstance(path, str) else path + + def _extract_from_obj(obj: Any, tokens: list[str], pos: int) -> list[str]: + if pos >= len(tokens): + if isinstance(obj, (str, int, float, bool)): + return [str(obj)] + elif obj is None: + return [] + elif isinstance(obj, (list, dict)): + return [json.dumps(obj, sort_keys=True)] + return [] + + token = tokens[pos] + results = [] + + if token.startswith("[") and token.endswith("]"): + if not isinstance(obj, list): + return [] + + index = token[1:-1] + if index == "*": + for item in obj: + results.extend(_extract_from_obj(item, tokens, pos + 1)) + else: + try: + idx = int(index) + if idx < 0: + idx = len(obj) + idx + if 0 <= idx < len(obj): + results.extend(_extract_from_obj(obj[idx], tokens, pos + 1)) + except (ValueError, IndexError): + return [] + + elif token.startswith("{") and token.endswith("}"): + if not isinstance(obj, dict): + return [] + + fields = [f.strip() for f in token[1:-1].split(",")] + for field in fields: + nested_tokens = tokenize_path(field) + if nested_tokens: + current_obj: Optional[dict] = obj + for nested_token in nested_tokens: + if ( + isinstance(current_obj, dict) + and nested_token in current_obj + ): + current_obj = current_obj[nested_token] + else: + current_obj = None + break + if current_obj is not None: + if isinstance(current_obj, (str, int, float, bool)): + results.append(str(current_obj)) + elif isinstance(current_obj, (list, dict)): + results.append(json.dumps(current_obj, sort_keys=True)) + + # Handle wildcard + elif token == "*": + if isinstance(obj, dict): + for value in obj.values(): + results.extend(_extract_from_obj(value, tokens, pos + 1)) + elif isinstance(obj, list): + for item in obj: + results.extend(_extract_from_obj(item, tokens, pos + 1)) + + # Handle regular field + else: + if isinstance(obj, dict) and token in obj: + results.extend(_extract_from_obj(obj[token], tokens, pos + 1)) + + return results + + return _extract_from_obj(obj, tokens, 0) + + +# Private utility functions + + +def tokenize_path(path: str) -> list[str]: + """Tokenize a path into components. + + !!! info "Types handled" + - Simple paths: "field1.field2" + - Array indexing: "[0]", "[*]", "[-1]" + - Wildcards: "*" + - Multi-field selection: "{field1,field2}" + """ + if not path: + return [] + + tokens = [] + current: list[str] = [] + i = 0 + while i < len(path): + char = path[i] + + if char == "[": # Handle array index + if current: + tokens.append("".join(current)) + current = [] + bracket_count = 1 + index_chars = ["["] + i += 1 + while i < len(path) and bracket_count > 0: + if path[i] == "[": + bracket_count += 1 + elif path[i] == "]": + bracket_count -= 1 + index_chars.append(path[i]) + i += 1 + tokens.append("".join(index_chars)) + continue + + elif char == "{": # Handle multi-field selection + if current: + tokens.append("".join(current)) + current = [] + brace_count = 1 + field_chars = ["{"] + i += 1 + while i < len(path) and brace_count > 0: + if path[i] == "{": + brace_count += 1 + elif path[i] == "}": + brace_count -= 1 + field_chars.append(path[i]) + i += 1 + tokens.append("".join(field_chars)) + continue + + elif char == ".": # Handle regular field + if current: + tokens.append("".join(current)) + current = [] + else: + current.append(char) + i += 1 + + if current: + tokens.append("".join(current)) + + return tokens + + +def _is_async_callable( + func: Any, +) -> bool: + """Check if a function is async. + + This includes both async def functions and classes with async __call__ methods. + + Args: + func: Function or callable object to check. + + Returns: + True if the function is async, False otherwise. + """ + return ( + asyncio.iscoroutinefunction(func) + or hasattr(func, "__call__") # noqa: B004 + and asyncio.iscoroutinefunction(func.__call__) + ) + + +__all__ = [ + "ensure_embeddings", + "EmbeddingsFunc", + "AEmbeddingsFunc", +] diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py index 69a315096..40011a7ba 100644 --- a/libs/checkpoint/langgraph/store/memory/__init__.py +++ b/libs/checkpoint/langgraph/store/memory/__init__.py @@ -1,79 +1,379 @@ +"""In-memory key-value store. + +A lightweight store implementation using Python dictionaries. Supports basic +key-value operations and vector search when configured with embeddings. + +Examples: + Basic key-value storage: + store = InMemoryStore() + store.put(("users", "123"), "prefs", {"theme": "dark"}) + item = store.get(("users", "123"), "prefs") + + Vector search with embeddings: + from langchain_openai import OpenAIEmbeddings + store = InMemoryStore(index={ + "dims": 1536, + "embed": OpenAIEmbeddings(model="text-embedding-3-small"), + }) + + # Store documents + store.put(("docs",), "doc1", {"text": "Python tutorial"}) + store.put(("docs",), "doc2", {"text": "TypeScript guide"}) + + # Search by similarity + results = store.search(("docs",), query="python programming") + + +Note: + For production use cases requiring persistence, use a database-backed store instead. +""" + +import asyncio +import concurrent.futures as cf +import functools +import logging from collections import defaultdict from datetime import datetime, timezone -from typing import Iterable +from importlib import util +from typing import Any, Iterable, Optional + +from langchain_core.embeddings import Embeddings from langgraph.store.base import ( BaseStore, GetOp, + IndexConfig, Item, ListNamespacesOp, MatchCondition, Op, PutOp, Result, + SearchItem, SearchOp, + ensure_embeddings, + get_text_at_path, + tokenize_path, ) +logger = logging.getLogger(__name__) + class InMemoryStore(BaseStore): - """A KV store backed by an in-memory python dictionary. + """In-memory dictionary-backed store with optional vector search. - Useful for testing/experimentation and lightweight PoC's. - For actual persistence, use a Store backed by a proper database. + Examples: + Basic key-value storage: + store = InMemoryStore() + store.put(("users", "123"), "prefs", {"theme": "dark"}) + item = store.get(("users", "123"), "prefs") + + Vector search with embeddings: + from langchain_openai import OpenAIEmbeddings + store = InMemoryStore(index={ + "dims": 1536, + "embed": OpenAIEmbeddings(model="text-embedding-3-small"), + }) + + # Store documents + store.put(("docs",), "doc1", {"text": "Python tutorial"}) + store.put(("docs",), "doc2", {"text": "TypeScript guide"}) + + # Search by similarity + results = store.search(("docs",), query="python programming") + + Warning: + This store keeps all data in memory. Data is lost when the process exits. + For persistence, use a database-backed store like PostgresStore. + + Tip: + For vector search, install numpy for better performance: + ```bash + pip install numpy + ``` """ - __slots__ = ("_data",) + __slots__ = ( + "_data", + "_vectors", + "index_config", + "embeddings", + ) - def __init__(self) -> 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._vectors: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = ( + defaultdict(lambda: defaultdict(dict)) + ) + self.index_config = index + if self.index_config: + self.index_config = self.index_config.copy() + self.embeddings: Optional[Embeddings] = ensure_embeddings( + self.index_config.get("embed"), + ) + 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.index_config = None + self.embeddings = None def batch(self, ops: Iterable[Op]) -> list[Result]: + # The batch/abatch methods are treated as internal. + # Users should access via put/search/get/list_namespaces/etc. + results, put_ops, search_ops = self._prepare_ops(ops) + if search_ops: + queryinmem_store = self._embed_search_queries(search_ops) + self._batch_search(search_ops, queryinmem_store, results) + + to_embed = self._extract_texts(put_ops) + 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) + return results + + async def abatch(self, ops: Iterable[Op]) -> list[Result]: + # The batch/abatch methods are treated as internal. + # Users should access via put/search/get/list_namespaces/etc. + results, put_ops, search_ops = self._prepare_ops(ops) + if search_ops: + queryinmem_store = await self._aembed_search_queries(search_ops) + self._batch_search(search_ops, queryinmem_store, results) + + to_embed = self._extract_texts(put_ops) + 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) + return results + + # Helpers + + def _filter_items(self, op: SearchOp) -> list[tuple[Item, list[list[float]]]]: + """Filter items by namespace and filter function, return items with their embeddings.""" + namespace_prefix = op.namespace_prefix + + def filter_func(item: Item) -> bool: + if not op.filter: + return True + + return all( + _compare_values(item.value.get(key), filter_value) + for key, filter_value in op.filter.items() + ) + + filtered = [] + for namespace in self._data: + if not ( + namespace[: len(namespace_prefix)] == namespace_prefix + if len(namespace) >= len(namespace_prefix) + else False + ): + continue + + for key, item in self._data[namespace].items(): + if filter_func(item): + if op.query and (embeddings := self._vectors[namespace].get(key)): + filtered.append((item, list(embeddings.values()))) + else: + filtered.append((item, [])) + return filtered + + def _embed_search_queries( + self, + search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]], + ) -> dict[str, list[float]]: + queryinmem_store = {} + if self.index_config and self.embeddings and search_ops: + queries = {op.query for (op, _) in search_ops.values() if op.query} + + if queries: + with cf.ThreadPoolExecutor() as executor: + futures = { + q: executor.submit(self.embeddings.embed_query, q) + for q in list(queries) + } + for query, future in futures.items(): + queryinmem_store[query] = future.result() + + return queryinmem_store + + async def _aembed_search_queries( + self, + search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]], + ) -> dict[str, list[float]]: + queryinmem_store = {} + if self.index_config and self.embeddings and search_ops: + queries = {op.query for (op, _) in search_ops.values() if op.query} + + if queries: + coros = [self.embeddings.aembed_query(q) for q in list(queries)] + results = await asyncio.gather(*coros) + queryinmem_store = dict(zip(queries, results)) + + return queryinmem_store + + def _batch_search( + self, + ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]], + queryinmem_store: dict[str, list[float]], + results: list[Result], + ) -> None: + """Perform batch similarity search for multiple queries.""" + for i, (op, candidates) in ops.items(): + if not candidates: + results[i] = [] + continue + 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: list[tuple[Optional[float], Item]] = [] + for score, item in sorted_results: + key = (item.namespace, item.key) + if key in seen: + continue + ix = len(seen) + seen.add(key) + if ix >= op.offset + op.limit: + break + if ix < op.offset: + 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( + namespace=item.namespace, + key=item.key, + value=item.value, + created_at=item.created_at, + updated_at=item.updated_at, + score=float(score) if score is not None else None, + ) + for score, item in kept + ] + else: + results[i] = [ + SearchItem( + namespace=item.namespace, + key=item.key, + value=item.value, + created_at=item.created_at, + updated_at=item.updated_at, + ) + for (item, _) in candidates[op.offset : op.offset + op.limit] + ] + + def _prepare_ops( + self, ops: Iterable[Op] + ) -> tuple[ + list[Result], + dict[tuple[tuple[str, ...], str], PutOp], + dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]], + ]: results: list[Result] = [] - for op in ops: + put_ops: dict[tuple[tuple[str, ...], str], PutOp] = {} + search_ops: dict[ + int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]] + ] = {} + for i, op in enumerate(ops): if isinstance(op, GetOp): item = self._data[op.namespace].get(op.key) results.append(item) elif isinstance(op, SearchOp): - candidates = [ - item - for namespace, items in self._data.items() - if ( - namespace[: len(op.namespace_prefix)] == op.namespace_prefix - if len(namespace) >= len(op.namespace_prefix) - else False - ) - for item in items.values() - ] - if op.filter: - candidates = [ - item - for item in candidates - if item.value.items() >= op.filter.items() - ] - results.append(candidates[op.offset : op.offset + op.limit]) - elif isinstance(op, PutOp): - if op.value is None: - self._data[op.namespace].pop(op.key, None) - elif op.key in self._data[op.namespace]: - self._data[op.namespace][op.key].value = op.value - self._data[op.namespace][op.key].updated_at = datetime.now( - timezone.utc - ) - else: - self._data[op.namespace][op.key] = Item( - value=op.value, - key=op.key, - namespace=op.namespace, - created_at=datetime.now(timezone.utc), - updated_at=datetime.now(timezone.utc), - ) + search_ops[i] = (op, self._filter_items(op)) results.append(None) elif isinstance(op, ListNamespacesOp): results.append(self._handle_list_namespaces(op)) - return results + elif isinstance(op, PutOp): + put_ops[(op.namespace, op.key)] = op + results.append(None) + else: + raise ValueError(f"Unknown operation type: {type(op)}") - async def abatch(self, ops: Iterable[Op]) -> list[Result]: - return self.batch(ops) + return results, put_ops, search_ops + + def _apply_put_ops(self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]) -> None: + for (namespace, key), op in put_ops.items(): + if op.value is None: + self._data[namespace].pop(key, None) + self._vectors[namespace].pop(key, None) + else: + self._data[namespace][key] = Item( + value=op.value, + key=key, + namespace=namespace, + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + 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.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: + 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: + for i, text in enumerate(texts): + to_embed[text].append( + (op.namespace, op.key, f"{path}.{i}") + ) + + else: + to_embed[texts[0]].append((op.namespace, op.key, path)) + + return to_embed + + return {} + + def _insertinmem_store( + self, + to_embed: dict[str, list[tuple[tuple[str, ...], str, str]]], + embeddings: list[list[float]], + ) -> None: + indices = [index for indices in to_embed.values() for index in indices] + if len(indices) != len(embeddings): + raise ValueError( + f"Number of embeddings ({len(embeddings)}) does not" + f" match number of indices ({len(indices)})" + ) + for embedding, (ns, key, path) in zip(embeddings, indices): + self._vectors[ns][key][path] = embedding def _handle_list_namespaces(self, op: ListNamespacesOp) -> list[tuple[str, ...]]: all_namespaces = list( @@ -94,7 +394,52 @@ class InMemoryStore(BaseStore): return namespaces[op.offset : op.offset + op.limit] +@functools.lru_cache(maxsize=1) +def _check_numpy() -> bool: + if bool(util.find_spec("numpy")): + return True + logger.warning( + "NumPy not found in the current Python environment. " + "The InMemoryStore will use a pure Python implementation for vector operations, " + "which may significantly impact performance, especially for large datasets or frequent searches. " + "For optimal speed and efficiency, consider installing NumPy: " + "pip install numpy" + ) + return False + + +def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]: + """ + Compute cosine similarity between a vector X and a matrix Y. + Lazy import numpy for efficiency. + """ + if _check_numpy(): + import numpy as np # type: ignore + + X_arr = np.array(X) if not isinstance(X, np.ndarray) else X + Y_arr = np.array(Y) if not isinstance(Y, np.ndarray) else Y + X_norm = np.linalg.norm(X_arr) + Y_norm = np.linalg.norm(Y_arr, axis=1) + + # Avoid division by zero + mask = Y_norm != 0 + similarities = np.zeros_like(Y_norm) + similarities[mask] = np.dot(Y_arr[mask], X_arr) / (Y_norm[mask] * X_norm) + return similarities.tolist() + + similarities = [] + for y in Y: + dot_product = sum(a * b for a, b in zip(X, y)) + norm1 = sum(a * a for a in X) ** 0.5 + norm2 = sum(a * a for a in y) ** 0.5 + similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0 + similarities.append(similarity) + + return similarities + + def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool: + """Whether a namespace key matches a match condition.""" match_type = match_condition.match_type path = match_condition.path @@ -117,3 +462,44 @@ def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool: return True else: raise ValueError(f"Unsupported match type: {match_type}") + + +def _compare_values(item_value: Any, filter_value: Any) -> bool: + """Compare values in a JSONB-like way, handling nested objects.""" + if isinstance(filter_value, dict): + if any(k.startswith("$") for k in filter_value): + return all( + _apply_operator(item_value, op_key, op_value) + for op_key, op_value in filter_value.items() + ) + if not isinstance(item_value, dict): + return False + return all( + _compare_values(item_value.get(k), v) for k, v in filter_value.items() + ) + elif isinstance(filter_value, (list, tuple)): + return ( + isinstance(item_value, (list, tuple)) + and len(item_value) == len(filter_value) + and all(_compare_values(iv, fv) for iv, fv in zip(item_value, filter_value)) + ) + else: + return item_value == filter_value + + +def _apply_operator(value: Any, operator: str, op_value: Any) -> bool: + """Apply a comparison operator, matching PostgreSQL's JSONB behavior.""" + if operator == "$eq": + return value == op_value + elif operator == "$gt": + return float(value) > float(op_value) + elif operator == "$gte": + return float(value) >= float(op_value) + elif operator == "$lt": + return float(value) < float(op_value) + elif operator == "$lte": + return float(value) <= float(op_value) + elif operator == "$ne": + return value != op_value + else: + raise ValueError(f"Unsupported operator: {operator}") diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index 278594fcb..deb7de5c4 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint" -version = "2.0.6" +version = "2.0.5" description = "Library with base interfaces for LangGraph checkpoint savers." authors = [] license = "MIT" diff --git a/libs/checkpoint/tests/embed_test_utils.py b/libs/checkpoint/tests/embed_test_utils.py new file mode 100644 index 000000000..d28cd959f --- /dev/null +++ b/libs/checkpoint/tests/embed_test_utils.py @@ -0,0 +1,55 @@ +"""Embedding utilities for testing.""" + +import math +import random +from collections import Counter, defaultdict +from typing import Any + +from langchain_core.embeddings import Embeddings + + +class CharacterEmbeddings(Embeddings): + """Simple character-frequency based embeddings using random projections.""" + + def __init__(self, dims: int = 50, seed: int = 42): + """Initialize with embedding dimensions and random seed.""" + self._rng = random.Random(seed) + self.dims = dims + # 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) + total = sum(counts.values()) + + if total == 0: + return [0.0] * self.dims + + embedding = [0.0] * self.dims + for char, count in counts.items(): + 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: + embedding = [x / norm for x in embedding] + + return embedding + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + """Embed a list of documents.""" + return [self._embed_one(text) for text in texts] + + def embed_query(self, text: str) -> list[float]: + """Embed a query string.""" + return self._embed_one(text) + + def __eq__(self, other: Any) -> bool: + return isinstance(other, CharacterEmbeddings) and self.dims == other.dims diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index 0ecd4bd84..053d9a981 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -1,19 +1,30 @@ +# mypy: disable-error-code="operator" import asyncio +import json from datetime import datetime -from typing import Iterable +from typing import Any, Iterable import pytest from pytest_mock import MockerFixture -from langgraph.store.base import GetOp, InvalidNamespaceError, Item, Op, PutOp, Result +from langgraph.store.base import ( + GetOp, + InvalidNamespaceError, + Item, + Op, + PutOp, + Result, + get_text_at_path, +) from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.memory import InMemoryStore +from tests.embed_test_utils import CharacterEmbeddings class MockAsyncBatchedStore(AsyncBatchedBaseStore): - def __init__(self) -> None: + def __init__(self, **kwargs: Any) -> None: super().__init__() - self._store = InMemoryStore() + self._store = InMemoryStore(**kwargs) def batch(self, ops: Iterable[Op]) -> list[Result]: return self._store.batch(ops) @@ -22,6 +33,74 @@ class MockAsyncBatchedStore(AsyncBatchedBaseStore): return self._store.batch(ops) +def test_get_text_at_path() -> None: + nested_data = { + "name": "test", + "info": { + "age": 25, + "tags": ["a", "b", "c"], + "metadata": {"created": "2024-01-01", "updated": "2024-01-02"}, + }, + "items": [ + {"id": 1, "value": "first", "tags": ["x", "y"]}, + {"id": 2, "value": "second", "tags": ["y", "z"]}, + {"id": 3, "value": "third", "tags": ["z", "w"]}, + ], + "empty": None, + "zeros": [0, 0.0, "0"], + "empty_list": [], + "empty_dict": {}, + } + + assert get_text_at_path(nested_data, "$") == [ + json.dumps(nested_data, sort_keys=True) + ] + + assert get_text_at_path(nested_data, "name") == ["test"] + assert get_text_at_path(nested_data, "info.age") == ["25"] + + assert get_text_at_path(nested_data, "info.metadata.created") == ["2024-01-01"] + + assert get_text_at_path(nested_data, "items[0].value") == ["first"] + assert get_text_at_path(nested_data, "items[-1].value") == ["third"] + assert get_text_at_path(nested_data, "items[1].tags[0]") == ["y"] + + values = get_text_at_path(nested_data, "items[*].value") + assert set(values) == {"first", "second", "third"} + + metadata_dates = get_text_at_path(nested_data, "info.metadata.*") + assert set(metadata_dates) == {"2024-01-01", "2024-01-02"} + name_and_age = get_text_at_path(nested_data, "{name,info.age}") + assert set(name_and_age) == {"test", "25"} + + item_fields = get_text_at_path(nested_data, "items[*].{id,value}") + assert set(item_fields) == {"1", "2", "3", "first", "second", "third"} + + all_tags = get_text_at_path(nested_data, "items[*].tags[*]") + assert set(all_tags) == {"x", "y", "z", "w"} + + assert get_text_at_path(None, "any.path") == [] + assert get_text_at_path({}, "any.path") == [] + assert get_text_at_path(nested_data, "") == [ + json.dumps(nested_data, sort_keys=True) + ] + assert get_text_at_path(nested_data, "nonexistent") == [] + assert get_text_at_path(nested_data, "items[99].value") == [] + assert get_text_at_path(nested_data, "items[*].nonexistent") == [] + + assert get_text_at_path(nested_data, "empty") == [] + assert get_text_at_path(nested_data, "empty_list") == ["[]"] + assert get_text_at_path(nested_data, "empty_dict") == ["{}"] + + zeros = get_text_at_path(nested_data, "zeros[*]") + assert set(zeros) == {"0", "0.0"} + + assert get_text_at_path(nested_data, "items[].value") == [] + assert get_text_at_path(nested_data, "items[abc].value") == [] + assert get_text_at_path(nested_data, "{unclosed") == [] + assert get_text_at_path(nested_data, "nested[{invalid}]") == [] + + async def test_async_batch_store(mocker: MockerFixture) -> None: abatch = mocker.stub() @@ -304,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 @@ -345,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 @@ -420,3 +504,446 @@ async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None: assert results[0][0].value == doc2 abatch.reset_mock() + + +@pytest.fixture +def fake_embeddings() -> CharacterEmbeddings: + return CharacterEmbeddings(dims=500) + + +def test_vector_store_initialization(fake_embeddings: CharacterEmbeddings) -> None: + """Test store initialization with embedding config.""" + store = InMemoryStore( + index={"dims": fake_embeddings.dims, "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( + fake_embeddings: CharacterEmbeddings, +) -> None: + """Test inserting items that get auto-embedded.""" + store = InMemoryStore( + index={"dims": fake_embeddings.dims, "embed": fake_embeddings} + ) + docs = [ + ("doc1", {"text": "short text"}), + ("doc2", {"text": "longer text document"}), + ("doc3", {"text": "longest text document here"}), + ("doc4", {"description": "text in description field"}), + ("doc5", {"content": "text in content field"}), + ("doc6", {"body": "text in body field"}), + ] + + for key, value in docs: + store.put(("test",), key, value) + + results = store.search(("test",), query="long text") + assert len(results) > 0 + + doc_order = [r.key for r in results] + assert "doc2" in doc_order + assert "doc3" in doc_order + + +async def test_async_vector_insert_with_auto_embedding( + fake_embeddings: CharacterEmbeddings, +) -> None: + """Test inserting items that get auto-embedded using async methods.""" + store = InMemoryStore( + index={"dims": fake_embeddings.dims, "embed": fake_embeddings} + ) + docs = [ + ("doc1", {"text": "short text"}), + ("doc2", {"text": "longer text document"}), + ("doc3", {"text": "longest text document here"}), + ("doc4", {"description": "text in description field"}), + ("doc5", {"content": "text in content field"}), + ("doc6", {"body": "text in body field"}), + ] + + for key, value in docs: + await store.aput(("test",), key, value) + + results = await store.asearch(("test",), query="long text") + assert len(results) > 0 + + doc_order = [r.key for r in results] + assert "doc2" in doc_order + assert "doc3" in doc_order + + +def test_vector_update_with_embedding(fake_embeddings: CharacterEmbeddings) -> None: + """Test that updating items properly updates their embeddings.""" + store = InMemoryStore( + index={"dims": fake_embeddings.dims, "embed": fake_embeddings} + ) + store.put(("test",), "doc1", {"text": "zany zebra Xerxes"}) + store.put(("test",), "doc2", {"text": "something about dogs"}) + store.put(("test",), "doc3", {"text": "text about birds"}) + + results_initial = store.search(("test",), query="Zany Xerxes") + assert len(results_initial) > 0 + assert results_initial[0].key == "doc1" + 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.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.score > after_score + + # Don't index this one + store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False) + results_new = store.search(("test",), query="new text about dogs", limit=3) + assert not any(r.key == "doc4" for r in results_new) + + +async def test_async_vector_update_with_embedding( + fake_embeddings: CharacterEmbeddings, +) -> None: + """Test that updating items properly updates their embeddings using async methods.""" + store = InMemoryStore( + 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"}) + await store.aput(("test",), "doc3", {"text": "text about birds"}) + + 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].score + + await store.aput(("test",), "doc1", {"text": "new text about dogs"}) + + results_after = await store.asearch(("test",), query="Zany Xerxes") + 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.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) + results_new = await store.asearch(("test",), query="new text about dogs", limit=3) + assert not any(r.key == "doc4" for r in results_new) + + +def test_vector_search_with_filters(fake_embeddings: CharacterEmbeddings) -> None: + """Test combining vector search with filters.""" + inmem_store = InMemoryStore( + index={"dims": fake_embeddings.dims, "embed": fake_embeddings} + ) + # Insert test documents + docs = [ + ("doc1", {"text": "red apple", "color": "red", "score": 4.5}), + ("doc2", {"text": "red car", "color": "red", "score": 3.0}), + ("doc3", {"text": "green apple", "color": "green", "score": 4.0}), + ("doc4", {"text": "blue car", "color": "blue", "score": 3.5}), + ] + + for key, value in docs: + inmem_store.put(("test",), key, value) + + results = inmem_store.search(("test",), query="apple", filter={"color": "red"}) + assert len(results) == 2 + assert results[0].key == "doc1" + + results = inmem_store.search(("test",), query="car", filter={"color": "red"}) + assert len(results) == 2 + assert results[0].key == "doc2" + + results = inmem_store.search( + ("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}} + ) + assert len(results) == 3 + assert results[0].key == "doc4" + + # Multiple filters + results = inmem_store.search( + ("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"} + ) + assert len(results) == 1 + assert results[0].key == "doc3" + + +async def test_async_vector_search_with_filters( + fake_embeddings: CharacterEmbeddings, +) -> None: + """Test combining vector search with filters using async methods.""" + store = InMemoryStore( + index={"dims": fake_embeddings.dims, "embed": fake_embeddings} + ) + # Insert test documents + docs = [ + ("doc1", {"text": "red apple", "color": "red", "score": 4.5}), + ("doc2", {"text": "red car", "color": "red", "score": 3.0}), + ("doc3", {"text": "green apple", "color": "green", "score": 4.0}), + ("doc4", {"text": "blue car", "color": "blue", "score": 3.5}), + ] + + for key, value in docs: + await store.aput(("test",), key, value) + + results = await store.asearch(("test",), query="apple", filter={"color": "red"}) + assert len(results) == 2 + assert results[0].key == "doc1" + + results = await store.asearch(("test",), query="car", filter={"color": "red"}) + assert len(results) == 2 + assert results[0].key == "doc2" + + results = await store.asearch( + ("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}} + ) + assert len(results) == 3 + assert results[0].key == "doc4" + + # Multiple filters + results = await store.asearch( + ("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"} + ) + assert len(results) == 1 + assert results[0].key == "doc3" + + +async def test_async_batched_vector_search_concurrent( + fake_embeddings: CharacterEmbeddings, +) -> None: + """Test concurrent vector search operations using async batched store.""" + store = MockAsyncBatchedStore( + index={"dims": fake_embeddings.dims, "embed": fake_embeddings} + ) + + colors = ["red", "blue", "green", "yellow", "purple"] + items = ["apple", "car", "house", "book", "phone"] + scores = [3.0, 3.5, 4.0, 4.5, 5.0] + + docs = [] + for i in range(50): + color = colors[i % len(colors)] + item = items[i % len(items)] + score = scores[i % len(scores)] + docs.append( + ( + f"doc{i}", + {"text": f"{color} {item}", "color": color, "score": score, "index": i}, + ) + ) + coros = [ + *[store.aput(("test",), key, value) for key, value in docs], + *[store.adelete(("test",), key) for key, value in docs], + *[store.aput(("test",), key, value) for key, value in docs], + ] + await asyncio.gather(*coros) + + # Prepare multiple search queries with different filters + search_queries: list[tuple[str, dict[str, Any]]] = [ + ("apple", {"color": "red"}), + ("car", {"color": "blue"}), + ("house", {"color": "green"}), + ("phone", {"score": {"$gt": 4.99}}), + ("book", {"score": {"$lte": 3.5}}), + ("apple", {"score": {"$gte": 3.0}, "color": "red"}), + ("car", {"score": {"$lt": 5.1}, "color": "blue"}), + ("house", {"index": {"$gt": 25}}), + ("phone", {"index": {"$lte": 10}}), + ] + + all_results = await asyncio.gather( + *[ + store.asearch(("test",), query=query, filter=filter_) + for query, filter_ in search_queries + ] + ) + + for results, (query, filter_) in zip(all_results, search_queries): + assert len(results) > 0, f"No results for query '{query}' with filter {filter_}" + + for result in results: + if "color" in filter_: + assert result.value["color"] == filter_["color"] + + if "score" in filter_: + score = result.value["score"] + for op, value in filter_["score"].items(): + if op == "$gt": + assert score > value + elif op == "$gte": + assert score >= value + elif op == "$lt": + assert score < value + elif op == "$lte": + assert score <= value + + if "index" in filter_: + index = result.value["index"] + for op, value in filter_["index"].items(): + if op == "$gt": + assert index > value + elif op == "$gte": + assert index >= value + elif op == "$lt": + assert index < value + elif op == "$lte": + assert index <= value + + +def test_vector_search_pagination(fake_embeddings: CharacterEmbeddings) -> None: + """Test pagination with vector search.""" + store = InMemoryStore( + 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}"}) + + results_page1 = store.search(("test",), query="test", limit=2) + results_page2 = store.search(("test",), query="test", limit=2, offset=2) + + assert len(results_page1) == 2 + assert len(results_page2) == 2 + assert results_page1[0].key != results_page2[0].key + + all_results = store.search(("test",), query="test", limit=10) + assert len(all_results) == 5 + + +async def test_async_vector_search_pagination( + fake_embeddings: CharacterEmbeddings, +) -> None: + """Test pagination with vector search using async methods.""" + store = InMemoryStore( + 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}"}) + + results_page1 = await store.asearch(("test",), query="test", limit=2) + results_page2 = await store.asearch(("test",), query="test", limit=2, offset=2) + + assert len(results_page1) == 2 + assert len(results_page2) == 2 + assert results_page1[0].key != results_page2[0].key + + all_results = await store.asearch(("test",), query="test", limit=10) + assert len(all_results) == 5 + + +async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None: + # Test store-level field configuration + store = InMemoryStore( + index={ + "dims": fake_embeddings.dims, + "embed": fake_embeddings, + # Key 2 isn't included. Don't index it. + "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].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].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].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