diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 707442d34..30276e47e 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -26,6 +26,7 @@ from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.types import SendProtocol +from langgraph.store.base import Item LC_REVIVER = Reviver() @@ -403,6 +404,18 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]: ), ), ) + elif isinstance(obj, Item): + return msgpack.ExtType( + EXT_CONSTRUCTOR_KW_ARGS, + _msgpack_enc( + ( + obj.__class__.__module__, + obj.__class__.__name__, + {k: getattr(obj, k) for k in obj.__slots__}, + ), + ), + ) + elif isinstance(obj, BaseException): return repr(obj) else: diff --git a/libs/checkpoint/langgraph/store/base.py b/libs/checkpoint/langgraph/store/base.py index 4f0aac7c9..c4b6bd552 100644 --- a/libs/checkpoint/langgraph/store/base.py +++ b/libs/checkpoint/langgraph/store/base.py @@ -5,55 +5,61 @@ """ from abc import ABC, abstractmethod -from dataclasses import dataclass from datetime import datetime -from typing import Any, Iterable, NamedTuple, Optional, Union +from typing import Any, Iterable, Literal, NamedTuple, Optional, Union -@dataclass class Item: - """Represents a stored item with metadata.""" - - value: dict[str, Any] - """The stored data as a dictionary. - - Keys are filterable. + """Represents a stored item with metadata. + + Args: + value (dict[str, Any]): The stored data as a dictionary. Keys are filterable. + (str): Unique identifier within the namespace. + namespace (tuple[str, ...]): Hierarchical path defining the collection in which this document resides. + Represented as a tuple of strings, allowing for nested categorization. + For example: ("documents", 'user123') + created_at (datetime): Timestamp of item creation. + updated_at (datetime): Timestamp of last update. """ - scores: dict[str, float] - """Relevance scores for the item. - - Keys can include built-in scores like 'recency' and 'relevance', - as well as any key present in the 'value' dictionary. This allows - for multi-dimensional scoring of items. - """ - - id: str - """Unique identifier within the namespace.""" - - namespace: tuple[str, ...] - """Hierarchical path defining the collection in which this document resides. - - Represented as a tuple of strings, allowing for nested categorization. - For example: ("documents", 'user123') - """ - - created_at: datetime - """Timestamp of item creation.""" - - updated_at: datetime - """Timestamp of last update.""" + __slots__ = ("value", "key", "namespace", "created_at", "updated_at") - last_accessed_at: datetime - """Timestamp of last access.""" + def __init__( + self, + *, + value: dict[str, Any], + key: str, + namespace: tuple[str, ...], + created_at: datetime, + updated_at: datetime, + ): + self.value = value + self.key = key + self.namespace = tuple(namespace) + self.created_at = created_at + self.updated_at = updated_at + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Item): + return False + return ( + self.value == other.value + and self.key == other.key + and self.namespace == other.namespace + and self.created_at == other.created_at + and self.updated_at == other.updated_at + ) + + def __hash__(self) -> int: + return hash((self.namespace, self.key)) class GetOp(NamedTuple): - """Operation to retrieve an item by namespace and ID.""" + """Operation to retrieve an item by namespace and key.""" namespace: tuple[str, ...] """Hierarchical path for the item.""" - id: str + key: str """Unique identifier within the namespace.""" @@ -80,7 +86,7 @@ class PutOp(NamedTuple): For example: ("documents", "user123") """ - id: str + key: str """Unique identifier for the document. Should be distinct within its namespace. @@ -97,8 +103,48 @@ class PutOp(NamedTuple): """ -Op = Union[GetOp, SearchOp, PutOp] -Result = Union[Item, list[Item], None] +NameSpacePath = tuple[Union[str, Literal["*"]], ...] + +NamespaceMatchType = Literal["prefix", "suffix"] + + +class MatchCondition(NamedTuple): + """Represents a single match condition.""" + + match_type: NamespaceMatchType + path: NameSpacePath + + +class ListNamespacesOp(NamedTuple): + """Operation to list namespaces with optional match conditions.""" + + match_conditions: Optional[tuple[MatchCondition, ...]] = None + """A tuple of match conditions to apply to namespaces.""" + + max_depth: Optional[int] = None + """Return namespaces up to this depth in the hierarchy.""" + + limit: int = 100 + """Maximum number of namespaces to return.""" + + offset: int = 0 + """Number of namespaces to skip before returning results.""" + + +Op = Union[GetOp, SearchOp, PutOp, ListNamespacesOp] +Result = Union[Item, list[Item], list[tuple[str, ...]], None] + + +class InvalidNamespaceError(ValueError): + """Provided namespace is invalid.""" + + +def _validate_namespace(namespace: tuple[str, ...]) -> None: + for label in namespace: + if "." in label: + raise InvalidNamespaceError( + f"Invalid namespace label '{label}'. Namespace labels cannot contain periods ('.')." + ) class BaseStore(ABC): @@ -114,17 +160,17 @@ def batch(self, ops: Iterable[Op]) -> list[Result]: async def abatch(self, ops: Iterable[Op]) -> list[Result]: """Execute a batch of operations asynchronously.""" - def get(self, namespace: tuple[str, ...], id: str) -> Optional[Item]: + def get(self, namespace: tuple[str, ...], key: str) -> Optional[Item]: """Retrieve a single item. Args: namespace: Hierarchical path for the item. - id: Unique identifier within the namespace. + key: Unique identifier within the namespace. Returns: The retrieved item or None if not found. """ - return self.batch([GetOp(namespace, id)])[0] + return self.batch([GetOp(namespace, key)])[0] def search( self, @@ -148,36 +194,88 @@ def search( """ return self.batch([SearchOp(namespace_prefix, filter, limit, offset)])[0] - def put(self, namespace: tuple[str, ...], id: str, value: dict[str, Any]) -> None: + def put(self, namespace: tuple[str, ...], key: str, value: dict[str, Any]) -> None: """Store or update an item. Args: namespace: Hierarchical path for the item. - id: Unique identifier within the namespace. + key: Unique identifier within the namespace. value: Dictionary containing the item's data. """ - self.batch([PutOp(namespace, id, value)]) + _validate_namespace(namespace) + self.batch([PutOp(namespace, key, value)]) - def delete(self, namespace: tuple[str, ...], id: str) -> None: + def delete(self, namespace: tuple[str, ...], key: str) -> None: """Delete an item. Args: namespace: Hierarchical path for the item. - id: Unique identifier within the namespace. + key: Unique identifier within the namespace. """ - self.batch([PutOp(namespace, id, None)]) + self.batch([PutOp(namespace, key, None)]) + + def list_namespaces( + self, + *, + prefix: Optional[NameSpacePath] = None, + suffix: Optional[NameSpacePath] = None, + max_depth: Optional[int] = None, + limit: int = 100, + offset: int = 0, + ) -> list[tuple[str, ...]]: + """List and filter namespaces in the store. + + Used to explore the organization of data, + find specific collections, or navigate the namespace hierarchy. + + Args: + 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. + limit (int): Maximum number of namespaces to return (default 100). + offset (int): Number of namespaces to skip for pagination (default 0). - async def aget(self, namespace: tuple[str, ...], id: str) -> Optional[Item]: + Returns: + List[Tuple[str, ...]]: A list of namespace tuples that match the criteria. + Each tuple represents a full namespace path up to `max_depth`. + + 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")] + """ + match_conditions = [] + if prefix: + match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) + if suffix: + match_conditions.append(MatchCondition(match_type="suffix", path=suffix)) + + op = ListNamespacesOp( + match_conditions=tuple(match_conditions), + max_depth=max_depth, + limit=limit, + offset=offset, + ) + return self.batch([op])[0] + + async def aget(self, namespace: tuple[str, ...], key: str) -> Optional[Item]: """Asynchronously retrieve a single item. Args: namespace: Hierarchical path for the item. - id: Unique identifier within the namespace. + key: Unique identifier within the namespace. Returns: The retrieved item or None if not found. """ - return (await self.abatch([GetOp(namespace, id)]))[0] + return (await self.abatch([GetOp(namespace, key)]))[0] async def asearch( self, @@ -204,22 +302,74 @@ async def asearch( ] async def aput( - self, namespace: tuple[str, ...], id: str, value: dict[str, Any] + self, namespace: tuple[str, ...], key: str, value: dict[str, Any] ) -> None: """Asynchronously store or update an item. Args: namespace: Hierarchical path for the item. - id: Unique identifier within the namespace. + key: Unique identifier within the namespace. value: Dictionary containing the item's data. """ - await self.abatch([PutOp(namespace, id, value)]) + _validate_namespace(namespace) + await self.abatch([PutOp(namespace, key, value)]) - async def adelete(self, namespace: tuple[str, ...], id: str) -> None: + async def adelete(self, namespace: tuple[str, ...], key: str) -> None: """Asynchronously delete an item. Args: namespace: Hierarchical path for the item. - id: Unique identifier within the namespace. + key: Unique identifier within the namespace. + """ + await self.abatch([PutOp(namespace, key, None)]) + + async def alist_namespaces( + self, + *, + prefix: Optional[NameSpacePath] = None, + suffix: Optional[NameSpacePath] = None, + max_depth: Optional[int] = None, + limit: int = 100, + offset: int = 0, + ) -> list[tuple[str, ...]]: + """List and filter namespaces in the store asynchronously. + + Used to explore the organization of data, + find specific collections, or navigate the namespace hierarchy. + + Args: + 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. + limit (int): Maximum number of namespaces to return (default 100). + offset (int): Number of namespaces to skip for pagination (default 0). + + Returns: + List[Tuple[str, ...]]: A list of namespace tuples that match the criteria. + Each tuple represents a full namespace path up to `max_depth`. + + 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") + await store.alist_namespaces(prefix=("a", "b"), max_depth=3) + # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] """ - await self.abatch([PutOp(namespace, id, None)]) + match_conditions = [] + if prefix: + match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) + if suffix: + match_conditions.append(MatchCondition(match_type="suffix", path=suffix)) + + op = ListNamespacesOp( + match_conditions=tuple(match_conditions), + max_depth=max_depth, + limit=limit, + offset=offset, + ) + return (await self.abatch([op]))[0] diff --git a/libs/checkpoint/langgraph/store/batch.py b/libs/checkpoint/langgraph/store/batch.py index 148acb8d1..8283a7a66 100644 --- a/libs/checkpoint/langgraph/store/batch.py +++ b/libs/checkpoint/langgraph/store/batch.py @@ -21,10 +21,10 @@ def __del__(self) -> None: async def aget( self, namespace: tuple[str, ...], - id: str, + key: str, ) -> Optional[Item]: fut = self._loop.create_future() - self._aqueue[fut] = GetOp(namespace, id) + self._aqueue[fut] = GetOp(namespace, key) return await fut async def asearch( @@ -32,9 +32,7 @@ async def asearch( namespace_prefix: tuple[str, ...], /, *, - query: Optional[str] = None, filter: Optional[dict[str, Any]] = None, - weights: Optional[dict[str, float]] = None, limit: int = 10, offset: int = 0, ) -> list[Item]: @@ -45,20 +43,20 @@ async def asearch( async def aput( self, namespace: tuple[str, ...], - id: str, + key: str, value: dict[str, Any], ) -> None: fut = self._loop.create_future() - self._aqueue[fut] = PutOp(namespace, id, value) + self._aqueue[fut] = PutOp(namespace, key, value) return await fut async def adelete( self, namespace: tuple[str, ...], - id: str, + key: str, ) -> None: fut = self._loop.create_future() - self._aqueue[fut] = PutOp(namespace, id, None) + self._aqueue[fut] = PutOp(namespace, key, None) return await fut diff --git a/libs/checkpoint/langgraph/store/memory.py b/libs/checkpoint/langgraph/store/memory.py index 2d002de9d..69a315096 100644 --- a/libs/checkpoint/langgraph/store/memory.py +++ b/libs/checkpoint/langgraph/store/memory.py @@ -2,7 +2,17 @@ from datetime import datetime, timezone from typing import Iterable -from langgraph.store.base import BaseStore, GetOp, Item, Op, PutOp, Result, SearchOp +from langgraph.store.base import ( + BaseStore, + GetOp, + Item, + ListNamespacesOp, + MatchCondition, + Op, + PutOp, + Result, + SearchOp, +) class InMemoryStore(BaseStore): @@ -21,9 +31,7 @@ def batch(self, ops: Iterable[Op]) -> list[Result]: results: list[Result] = [] for op in ops: if isinstance(op, GetOp): - item = self._data[op.namespace].get(op.id) - if item is not None: - item.last_accessed_at = datetime.now(timezone.utc) + item = self._data[op.namespace].get(op.key) results.append(item) elif isinstance(op, SearchOp): candidates = [ @@ -45,24 +53,67 @@ def batch(self, ops: Iterable[Op]) -> list[Result]: results.append(candidates[op.offset : op.offset + op.limit]) elif isinstance(op, PutOp): if op.value is None: - self._data[op.namespace].pop(op.id, None) - elif op.id in self._data[op.namespace]: - self._data[op.namespace][op.id].value = op.value - self._data[op.namespace][op.id].updated_at = datetime.now( + 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.id] = Item( + self._data[op.namespace][op.key] = Item( value=op.value, - scores={}, - id=op.id, + key=op.key, namespace=op.namespace, created_at=datetime.now(timezone.utc), updated_at=datetime.now(timezone.utc), - last_accessed_at=datetime.now(timezone.utc), ) results.append(None) + elif isinstance(op, ListNamespacesOp): + results.append(self._handle_list_namespaces(op)) return results async def abatch(self, ops: Iterable[Op]) -> list[Result]: return self.batch(ops) + + def _handle_list_namespaces(self, op: ListNamespacesOp) -> list[tuple[str, ...]]: + all_namespaces = list( + self._data.keys() + ) # Avoid collection size changing while iterating + namespaces = all_namespaces + if op.match_conditions: + namespaces = [ + ns + for ns in namespaces + if all(_does_match(condition, ns) for condition in op.match_conditions) + ] + + if op.max_depth is not None: + namespaces = sorted({ns[: op.max_depth] for ns in namespaces}) + else: + namespaces = sorted(namespaces) + return namespaces[op.offset : op.offset + op.limit] + + +def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool: + match_type = match_condition.match_type + path = match_condition.path + + if len(key) < len(path): + return False + + if match_type == "prefix": + for k_elem, p_elem in zip(key, path): + if p_elem == "*": + continue # Wildcard matches any element + if k_elem != p_elem: + return False + return True + elif match_type == "suffix": + for k_elem, p_elem in zip(reversed(key), reversed(path)): + if p_elem == "*": + continue # Wildcard matches any element + if k_elem != p_elem: + return False + return True + else: + raise ValueError(f"Unsupported match type: {match_type}") diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index e52531414..e504e80b1 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -15,6 +15,7 @@ from zoneinfo import ZoneInfo from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.store.base import Item class InnerPydantic(BaseModel): @@ -121,6 +122,13 @@ def test_serde_jsonplus() -> None: "a_float": 1.1, "a_bytes": b"my bytes", "a_bytearray": bytearray([42]), + "my_item": Item( + value={}, + key="my-key", + namespace=("a", "name", " "), + created_at=datetime(2024, 9, 24, 17, 29, 10, 128397), + updated_at=datetime(2024, 9, 24, 17, 29, 10, 128397), + ), } serde = JsonPlusSerializer() diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index fecc54dc1..317d9a720 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -6,6 +6,7 @@ from langgraph.store.base import GetOp, Item, Op, Result from langgraph.store.batch import AsyncBatchedBaseStore +from langgraph.store.memory import InMemoryStore async def test_async_batch_store(mocker: MockerFixture) -> None: @@ -21,12 +22,10 @@ async def abatch(self, ops: Iterable[Op]) -> list[Result]: return [ Item( value={}, - scores={}, - id=getattr(op, "id", ""), + key=getattr(op, "key", ""), namespace=getattr(op, "namespace", ()), created_at=datetime(2024, 9, 24, 17, 29, 10, 128397), updated_at=datetime(2024, 9, 24, 17, 29, 10, 128397), - last_accessed_at=datetime(2024, 9, 24, 17, 29, 10, 128397), ) for op in ops ] @@ -35,27 +34,23 @@ async def abatch(self, ops: Iterable[Op]) -> list[Result]: # concurrent calls are batched results = await asyncio.gather( - store.aget(namespace=("a",), id="b"), - store.aget(namespace=("c",), id="d"), + store.aget(namespace=("a",), key="b"), + store.aget(namespace=("c",), key="d"), ) assert results == [ Item( - {}, - {}, - "b", - ("a",), - datetime(2024, 9, 24, 17, 29, 10, 128397), - datetime(2024, 9, 24, 17, 29, 10, 128397), - datetime(2024, 9, 24, 17, 29, 10, 128397), + value={}, + key="b", + namespace=("a",), + created_at=datetime(2024, 9, 24, 17, 29, 10, 128397), + updated_at=datetime(2024, 9, 24, 17, 29, 10, 128397), ), Item( - {}, - {}, - "d", - ("c",), - datetime(2024, 9, 24, 17, 29, 10, 128397), - datetime(2024, 9, 24, 17, 29, 10, 128397), - datetime(2024, 9, 24, 17, 29, 10, 128397), + value={}, + key="d", + namespace=("c",), + created_at=datetime(2024, 9, 24, 17, 29, 10, 128397), + updated_at=datetime(2024, 9, 24, 17, 29, 10, 128397), ), ] assert abatch.call_count == 1 @@ -65,3 +60,202 @@ async def abatch(self, ops: Iterable[Op]) -> list[Result]: GetOp(("c",), "d"), ), ] + + +def test_list_namespaces_basic() -> None: + store = InMemoryStore() + + namespaces = [ + ("a", "b", "c"), + ("a", "b", "d", "e"), + ("a", "b", "d", "i"), + ("a", "b", "f"), + ("a", "c", "f"), + ("b", "a", "f"), + ("users", "123"), + ("users", "456", "settings"), + ("admin", "users", "789"), + ] + + for i, ns in enumerate(namespaces): + store.put(namespace=ns, key=f"id_{i}", value={"data": f"value_{i:02d}"}) + + result = store.list_namespaces(prefix=("a", "b")) + expected = [ + ("a", "b", "c"), + ("a", "b", "d", "e"), + ("a", "b", "d", "i"), + ("a", "b", "f"), + ] + assert sorted(result) == sorted(expected) + + result = store.list_namespaces(suffix=("f",)) + expected = [ + ("a", "b", "f"), + ("a", "c", "f"), + ("b", "a", "f"), + ] + assert sorted(result) == sorted(expected) + + result = store.list_namespaces(prefix=("a",), suffix=("f",)) + expected = [ + ("a", "b", "f"), + ("a", "c", "f"), + ] + assert sorted(result) == sorted(expected) + + # Test max_depth + result = store.list_namespaces(prefix=("a", "b"), max_depth=3) + expected = [ + ("a", "b", "c"), + ("a", "b", "d"), + ("a", "b", "f"), + ] + assert sorted(result) == sorted(expected) + + # Test limit and offset + result = store.list_namespaces(prefix=("a", "b"), limit=2) + expected = [ + ("a", "b", "c"), + ("a", "b", "d", "e"), + ] + assert result == expected + + result = store.list_namespaces(prefix=("a", "b"), offset=2) + expected = [ + ("a", "b", "d", "i"), + ("a", "b", "f"), + ] + assert result == expected + + result = store.list_namespaces(prefix=("a", "*", "f")) + expected = [ + ("a", "b", "f"), + ("a", "c", "f"), + ] + assert sorted(result) == sorted(expected) + + result = store.list_namespaces(suffix=("*", "f")) + expected = [ + ("a", "b", "f"), + ("a", "c", "f"), + ("b", "a", "f"), + ] + assert sorted(result) == sorted(expected) + + result = store.list_namespaces(prefix=("nonexistent",)) + assert result == [] + + result = store.list_namespaces(prefix=("users", "123")) + expected = [("users", "123")] + assert result == expected + + +def test_list_namespaces_with_wildcards() -> None: + store = InMemoryStore() + + namespaces = [ + ("users", "123"), + ("users", "456"), + ("users", "789", "settings"), + ("admin", "users", "789"), + ("guests", "123"), + ("guests", "456", "preferences"), + ] + + for i, ns in enumerate(namespaces): + store.put(namespace=ns, key=f"id_{i}", value={"data": f"value_{i:02d}"}) + + result = store.list_namespaces(prefix=("users", "*")) + expected = [ + ("users", "123"), + ("users", "456"), + ("users", "789", "settings"), + ] + assert sorted(result) == sorted(expected) + + result = store.list_namespaces(suffix=("*", "preferences")) + expected = [ + ("guests", "456", "preferences"), + ] + assert result == expected + + result = store.list_namespaces(prefix=("*", "users"), suffix=("*", "settings")) + + assert result == [] + + store.put( + namespace=("admin", "users", "settings", "789"), + key="foo", + value={"data": "some_val"}, + ) + expected = [ + ("admin", "users", "settings", "789"), + ] + + +def test_list_namespaces_pagination() -> None: + store = InMemoryStore() + + for i in range(20): + ns = ("namespace", f"sub_{i:02d}") + store.put(namespace=ns, key=f"id_{i:02d}", value={"data": f"value_{i:02d}"}) + + result = store.list_namespaces(prefix=("namespace",), limit=5, offset=0) + expected = [("namespace", f"sub_{i:02d}") for i in range(5)] + assert result == expected + + result = store.list_namespaces(prefix=("namespace",), limit=5, offset=5) + expected = [("namespace", f"sub_{i:02d}") for i in range(5, 10)] + assert result == expected + + result = store.list_namespaces(prefix=("namespace",), limit=5, offset=15) + expected = [("namespace", f"sub_{i:02d}") for i in range(15, 20)] + assert result == expected + + +def test_list_namespaces_max_depth() -> None: + store = InMemoryStore() + + namespaces = [ + ("a", "b", "c", "d"), + ("a", "b", "c", "e"), + ("a", "b", "f"), + ("a", "g"), + ("h", "i", "j", "k"), + ] + + for i, ns in enumerate(namespaces): + store.put(namespace=ns, key=f"id_{i}", value={"data": f"value_{i:02d}"}) + + result = store.list_namespaces(max_depth=2) + expected = [ + ("a", "b"), + ("a", "g"), + ("h", "i"), + ] + assert sorted(result) == sorted(expected) + + +def test_list_namespaces_no_conditions() -> None: + store = InMemoryStore() + + namespaces = [ + ("a", "b"), + ("c", "d"), + ("e", "f", "g"), + ] + + for i, ns in enumerate(namespaces): + store.put(namespace=ns, key=f"id_{i}", value={"data": f"value_{i:02d}"}) + + result = store.list_namespaces() + expected = namespaces + assert sorted(result) == sorted(expected) + + +def test_list_namespaces_empty_store() -> None: + store = InMemoryStore() + + result = store.list_namespaces() + assert result == [] diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py index 12a15fc93..7c8f45b14 100644 --- a/libs/langgraph/langgraph/managed/shared_value.py +++ b/libs/langgraph/langgraph/managed/shared_value.py @@ -59,7 +59,7 @@ def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]: with super().enter(config, **kwargs) as value: if value.store is not None: saved = value.store.search(value.ns) - value.value = {it.id: it.value for it in saved} + value.value = {it.key: it.value for it in saved} yield value @classmethod @@ -68,7 +68,7 @@ async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Se async with super().aenter(config, **kwargs) as value: if value.store is not None: saved = await value.store.asearch(value.ns) - value.value = {it.id: it.value for it in saved} + value.value = {it.key: it.value for it in saved} yield value def __init__(