Merge pull request #1352 from langchain-ai/nc/14aug/scoped-value

Add SharedValue

libs/langgraph/langgraph/constants.py

@@ -5,6 +5,7 @@ INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
+CONFIG_KEY_STORE = "__pregel_store"
 CONFIG_KEY_RESUMING = "__pregel_resuming"
 CONFIG_KEY_TASK_ID = "__pregel_task_id"
 INTERRUPT = "__interrupt__"
@@ -17,6 +18,7 @@ RESERVED = {
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
     CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_STORE,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,

libs/langgraph/langgraph/graph/state.py

@@ -40,10 +40,18 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
-from langgraph.managed.base import ManagedValue, is_managed_value
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ChannelTypePlaceholder,
+    ConfiguredManagedValue,
+    ManagedValue,
+    is_managed_value,
+    is_writable_managed_value,
+)
 from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 from langgraph.utils import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
@@ -374,6 +382,8 @@ class StateGraph(Graph):
     def compile(
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
+        *,
+        store: Optional[BaseStore] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
@@ -442,6 +452,7 @@ class StateGraph(Graph):
             interrupt_after_nodes=interrupt_after,
             auto_validate=False,
             debug=debug,
+            store=store,
         )
 
         compiled.attach_node(START, None)
@@ -511,7 +522,11 @@ class CompiledStateGraph(CompiledGraph):
                 if not isinstance(v, Context) and not is_managed_value(v)
             ]
         else:
-            output_keys = list(self.builder.channels)
+            output_keys = list(self.builder.channels) + [
+                k
+                for k, v in self.builder.managed.items()
+                if is_writable_managed_value(v)
+            ]
 
         def _get_state_key(
             input: Union[None, dict, Any], config: RunnableConfig, *, key: str
@@ -681,10 +696,10 @@ def _get_channels(
     schema: Type[dict],
 ) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]:
     if not hasattr(schema, "__annotations__"):
-        return {"__root__": _get_channel(schema, allow_managed=False)}, {}
+        return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
 
     all_keys = {
-        name: _get_channel(typ)
+        name: _get_channel(name, typ)
         for name, typ in get_type_hints(schema, include_extras=True).items()
         if name != "__slots__"
     }
@@ -695,9 +710,9 @@ def _get_channels(
 
 
 def _get_channel(
-    annotation: Any, *, allow_managed: bool = True
+    name: str, annotation: Any, *, allow_managed: bool = True
 ) -> Union[BaseChannel, Type[ManagedValue]]:
-    if manager := _is_field_managed_value(annotation):
+    if manager := _is_field_managed_value(name, annotation):
         if allow_managed:
             return manager
         else:
@@ -736,12 +751,18 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     return None
 
 
-def _is_field_managed_value(typ: Type[Any]) -> Optional[Type[ManagedValue]]:
+def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[Type[ManagedValue]]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1:
             decoration = get_origin(meta[-1]) or meta[-1]
             if is_managed_value(decoration):
+                if isinstance(decoration, ConfiguredManagedValue):
+                    for k, v in decoration.kwargs.items():
+                        if v is ChannelKeyPlaceholder:
+                            decoration.kwargs[k] = name
+                        if v is ChannelTypePlaceholder:
+                            decoration.kwargs[k] = typ.__origin__
                 return decoration
 
     return None

libs/langgraph/langgraph/managed/base.py

@@ -5,9 +5,12 @@ from inspect import isclass
 from typing import (
     Any,
     AsyncGenerator,
+    AsyncIterator,
     Generator,
     Generic,
+    Iterator,
     NamedTuple,
+    Sequence,
     Type,
     TypeVar,
     Union,
@@ -17,6 +20,7 @@ from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self, TypeGuard
 
 V = TypeVar("V")
+U = TypeVar("U")
 
 
 class ManagedValue(ABC, Generic[V]):
@@ -25,9 +29,7 @@ class ManagedValue(ABC, Generic[V]):
 
     @classmethod
     @contextmanager
-    def enter(
-        cls, config: RunnableConfig, **kwargs: Any
-    ) -> Generator[Self, None, None]:
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
         try:
             value = cls(config, **kwargs)
             yield value
@@ -41,9 +43,7 @@ class ManagedValue(ABC, Generic[V]):
 
     @classmethod
     @asynccontextmanager
-    async def aenter(
-        cls, config: RunnableConfig, **kwargs: Any
-    ) -> AsyncGenerator[Self, None]:
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
         try:
             value = cls(config, **kwargs)
             yield value
@@ -60,6 +60,16 @@ class ManagedValue(ABC, Generic[V]):
         ...
 
 
+class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC):
+    @abstractmethod
+    def update(self, writes: Sequence[U]) -> None:
+        ...
+
+    @abstractmethod
+    async def aupdate(self, writes: Sequence[U]) -> None:
+        ...
+
+
 class ConfiguredManagedValue(NamedTuple):
     cls: Type[ManagedValue]
     kwargs: dict[str, Any]
@@ -76,6 +86,24 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
     )
 
 
+def is_readonly_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]:
+    return (
+        isclass(value)
+        and issubclass(value, ManagedValue)
+        and not issubclass(value, WritableManagedValue)
+    ) or (
+        isinstance(value, ConfiguredManagedValue)
+        and not issubclass(value.cls, WritableManagedValue)
+    )
+
+
+def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue]]:
+    return (isclass(value) and issubclass(value, WritableManagedValue)) or (
+        isinstance(value, ConfiguredManagedValue)
+        and issubclass(value.cls, WritableManagedValue)
+    )
+
+
 @contextmanager
 def ManagedValuesManager(
     values: dict[str, ManagedValueSpec],
@@ -119,3 +147,7 @@ async def AsyncManagedValuesManager(
             yield {tasks[task]: task.result() for task in done}
     else:
         yield {}
+
+
+ChannelKeyPlaceholder = object()
+ChannelTypePlaceholder = object()

libs/langgraph/langgraph/managed/shared_value.py

@@ -0,0 +1,126 @@
+import collections.abc
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterator,
+    Optional,
+    Sequence,
+    Type,
+)
+
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import NotRequired, Required, Self
+
+from langgraph.constants import CONFIG_KEY_STORE
+from langgraph.errors import InvalidUpdateError
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ChannelTypePlaceholder,
+    ConfiguredManagedValue,
+    WritableManagedValue,
+)
+from langgraph.store.base import BaseStore
+
+V = dict[str, Any]
+
+
+Value = dict[str, V]
+Update = dict[str, Optional[V]]
+
+
+# Adapted from typing_extensions
+def _strip_extras(t):
+    """Strips Annotated, Required and NotRequired from a given type."""
+    if hasattr(t, "__origin__"):
+        return _strip_extras(t.__origin__)
+    if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
+        return _strip_extras(t.__args__[0])
+
+    return t
+
+
+class SharedValue(WritableManagedValue[Value, Update]):
+    @staticmethod
+    def on(scope: str) -> ConfiguredManagedValue:
+        return ConfiguredManagedValue(
+            SharedValue,
+            {
+                "scope": scope,
+                "key": ChannelKeyPlaceholder,
+                "typ": ChannelTypePlaceholder,
+            },
+        )
+
+    @classmethod
+    @contextmanager
+    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.list([value.ns])
+                value.value = saved[value.ns] or {}
+            yield value
+
+    @classmethod
+    @asynccontextmanager
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
+        async with super().aenter(config, **kwargs) as value:
+            if value.store is not None:
+                saved = await value.store.alist([value.ns])
+                value.value = saved[value.ns] or {}
+            yield value
+
+    def __init__(
+        self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
+    ) -> None:
+        if typ := _strip_extras(typ):
+            if typ not in (
+                dict,
+                collections.abc.Mapping,
+                collections.abc.MutableMapping,
+            ):
+                raise ValueError("SharedValue must be a dict")
+        self.scope = scope
+        self.config = config
+        self.value: Value = {}
+        self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE)
+        if self.store is None:
+            self.ns: Optional[str] = None
+        elif scope_value := config["configurable"].get(self.scope):
+            self.ns = f"scoped:{scope}:{key}:{scope_value}"
+        else:
+            raise ValueError(
+                f"Scope {scope} for shared state key not in config.configurable"
+            )
+
+    def __call__(self, step: int) -> Value:
+        return self.value.copy()
+
+    def _process_update(
+        self, values: Sequence[Update]
+    ) -> list[tuple[str, str, Optional[dict[str, Any]]]]:
+        writes = []
+        for vv in values:
+            for k, v in vv.items():
+                if v is None:
+                    if k in self.value:
+                        self.value[k] = None
+                        writes.append((self.ns, k, None))
+                elif not isinstance(v, dict):
+                    raise InvalidUpdateError("Received a non-dict value")
+                else:
+                    self.value[k] = v
+                    writes.append((self.ns, k, v))
+        return writes
+
+    def update(self, values: Sequence[Update]) -> None:
+        if self.store is None:
+            self._process_update(values)
+        else:
+            return self.store.put(self._process_update(values))
+
+    async def aupdate(self, writes: Sequence[Update]) -> None:
+        if self.store is None:
+            self._process_update(writes)
+        else:
+            return await self.store.aput(self._process_update(writes))

libs/langgraph/langgraph/pregel/__init__.py

@@ -108,6 +108,7 @@ from langgraph.pregel.types import (
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 
 WriteValue = Union[
     Runnable[Input, Output],
@@ -223,6 +224,9 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
+    store: Optional[BaseStore] = None
+    """Memory store to use for SharedValues. Defaults to None."""
+
     retry_policy: Optional[RetryPolicy] = None
     """Retry policy to use when running tasks. Set to None to disable."""
 
@@ -644,9 +648,9 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
-            )
+            ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
             # check interrupt before
             if tasks := should_interrupt(
@@ -788,9 +792,9 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
-            )
+            ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
             # check interrupt before
             if tasks := should_interrupt(

libs/langgraph/langgraph/pregel/algo.py

@@ -121,6 +121,7 @@ def local_write(
     commit: Callable[[Sequence[tuple[str, Any]]], None],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
     writes: Sequence[tuple[str, Any]],
 ) -> None:
     for chan, value in writes:
@@ -131,7 +132,7 @@ def local_write(
                 )
             if value.node not in processes:
                 raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
-        elif chan not in channels:
+        elif chan not in channels and chan not in managed:
             logger.warning(f"Skipping write for channel '{chan}' which has no readers")
     commit(writes)
 
@@ -145,7 +146,7 @@ def apply_writes(
     channels: Mapping[str, BaseChannel],
     tasks: Sequence[WritesProtocol],
     get_next_version: Optional[Callable[[int, BaseChannel], int]],
-) -> None:
+) -> dict[str, list[Any]]:
     # update seen versions
     for task in tasks:
         checkpoint["versions_seen"].setdefault(task.name, {}).update(
@@ -161,6 +162,7 @@ def apply_writes(
         max_version = max(checkpoint["channel_versions"].values())
     else:
         max_version = None
+
     # Consume all channels that were read
     for chan in {
         chan for task in tasks for chan in task.triggers if chan not in RESERVED
@@ -177,12 +179,15 @@ def apply_writes(
 
     # Group writes by channel
     pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
+    pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
             if chan == TASKS:
                 checkpoint["pending_sends"].append(val)
-            else:
+            elif chan in channels:
                 pending_writes_by_channel[chan].append(val)
+            else:
+                pending_writes_by_managed[chan].append(val)
 
     # Find the highest version of all channels
     if checkpoint["channel_versions"]:
@@ -214,6 +219,9 @@ def apply_writes(
                     max_version, channels[chan]
                 )
 
+    # Return managed values writes to be applied externally
+    return pending_writes_by_managed
+
 
 @overload
 def prepare_next_tasks(
@@ -314,7 +322,11 @@ def prepare_next_tasks(
                                 CONFIG_KEY_TASK_ID: task_id,
                                 # deque.extend is thread-safe
                                 CONFIG_KEY_SEND: partial(
-                                    local_write, writes.extend, processes, channels
+                                    local_write,
+                                    writes.extend,
+                                    processes,
+                                    channels,
+                                    managed,
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
@@ -405,7 +417,11 @@ def prepare_next_tasks(
                                     CONFIG_KEY_TASK_ID: task_id,
                                     # deque.extend is thread-safe
                                     CONFIG_KEY_SEND: partial(
-                                        local_write, writes.extend, processes, channels
+                                        local_write,
+                                        writes.extend,
+                                        processes,
+                                        channels,
+                                        managed,
                                     ),
                                     CONFIG_KEY_READ: partial(
                                         local_read,

libs/langgraph/langgraph/pregel/loop.py

@@ -18,10 +18,11 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
-from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables import RunnableConfig, patch_config
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel
@@ -42,6 +43,7 @@ from langgraph.checkpoint.base import (
 from langgraph.constants import (
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
+    CONFIG_KEY_STORE,
     ERROR,
     INPUT,
     INTERRUPT,
@@ -52,6 +54,7 @@ from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValueMapping,
     ManagedValuesManager,
+    WritableManagedValue,
 )
 from langgraph.pregel.algo import (
     PregelTaskWrites,
@@ -69,6 +72,8 @@ from langgraph.pregel.executor import (
 from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
 from langgraph.pregel.types import PregelExecutableTask
 from langgraph.pregel.utils import get_new_channel_versions
+from langgraph.store.base import BaseStore
+from langgraph.store.batch import AsyncBatchedStore
 
 if TYPE_CHECKING:
     from langgraph.pregel import Pregel
@@ -100,7 +105,7 @@ class PregelLoop:
         ]
     ]
     graph: "Pregel"
-
+    store: Optional[BaseStore]
     submit: Submit
     channels: Mapping[str, BaseChannel]
     managed: ManagedValueMapping
@@ -182,12 +187,15 @@ class PregelLoop:
         elif all(task.writes for task in self.tasks):
             writes = [w for t in self.tasks for w in t.writes]
             # all tasks have finished
-            apply_writes(
+            mv_writes = apply_writes(
                 self.checkpoint,
                 self.channels,
                 self.tasks,
                 self.checkpointer_get_next_version,
             )
+            # apply writes to managed values
+            for key, values in mv_writes.items():
+                self._update_mv(key, values)
             # produce values output
             self.stream.extend(
                 ("values", v)
@@ -324,12 +332,12 @@ class PregelLoop:
                 manager=None,
             )
             # apply input writes
-            apply_writes(
+            assert not apply_writes(
                 self.checkpoint,
                 self.channels,
                 discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
                 self.checkpointer_get_next_version,
-            )
+            ), "Can't write to SharedValues in graph input"
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": self.input})
         else:
@@ -395,6 +403,9 @@ class PregelLoop:
         # increment step
         self.step += 1
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        raise NotImplementedError
+
     def _suppress_interrupt(
         self,
         exc_type: Optional[Type[BaseException]],
@@ -415,6 +426,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         graph: "Pregel",
     ) -> None:
         super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        self.store = graph.store
         self.stack = ExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -438,6 +450,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         finally:
             self.checkpointer.put(config, checkpoint, metadata, new_versions)
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        return self.submit(cast(WritableManagedValue, self.managed[key]).update, values)
+
     # context manager
 
     def __enter__(self) -> Self:
@@ -461,7 +476,10 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             ChannelsManager(self.graph.channels, self.checkpoint, self.config)
         )
         self.managed = self.stack.enter_context(
-            ManagedValuesManager(self.graph.managed_values_dict, self.config)
+            ManagedValuesManager(
+                self.graph.managed_values_dict,
+                patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}),
+            )
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
@@ -492,6 +510,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         graph: "Pregel",
     ) -> None:
         super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        self.store = AsyncBatchedStore(graph.store) if graph.store else None
         self.stack = AsyncExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -515,6 +534,11 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         finally:
             await self.checkpointer.aput(config, checkpoint, metadata, new_versions)
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        return self.submit(
+            cast(WritableManagedValue, self.managed[key]).aupdate, values
+        )
+
     # context manager
 
     async def __aenter__(self) -> Self:
@@ -540,7 +564,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
         )
         self.managed = await self.stack.enter_async_context(
-            AsyncManagedValuesManager(self.graph.managed_values_dict, self.config)
+            AsyncManagedValuesManager(
+                self.graph.managed_values_dict,
+                patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}),
+            )
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"

libs/langgraph/langgraph/store/base.py

@@ -0,0 +1,21 @@
+from typing import Any, List, Optional
+
+V = dict[str, Any]
+
+
+class BaseStore:
+    def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        # list[namespace] -> dict[namespace, list[value]]
+        raise NotImplementedError
+
+    def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        # list[(namespace, key, value | none)] -> None
+        raise NotImplementedError
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        # list[namespace] -> dict[namespace, list[value]]
+        raise NotImplementedError
+
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        # list[(namespace, key, value | none)] -> None
+        raise NotImplementedError

libs/langgraph/langgraph/store/batch.py

@@ -0,0 +1,65 @@
+import asyncio
+from typing import NamedTuple, Optional, Union
+
+from langgraph.store.base import BaseStore, V
+
+
+class ListOp(NamedTuple):
+    prefixes: list[str]
+
+
+class PutOp(NamedTuple):
+    writes: list[tuple[str, str, Optional[V]]]
+
+
+class AsyncBatchedStore(BaseStore):
+    def __init__(self, store: BaseStore) -> None:
+        self.store = store
+        self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {}
+        self.task = asyncio.create_task(_run(self.aqueue, self.store))
+
+    def __del__(self) -> None:
+        self.task.cancel()
+
+    async def alist(self, prefixes: list[str]) -> dict[str, dict[str, V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = ListOp(prefixes)
+        return await fut
+
+    async def aput(self, writes: list[tuple[str, str, Optional[V]]]) -> None:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = PutOp(writes)
+        return await fut
+
+
+async def _run(
+    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], store: BaseStore
+) -> None:
+    while True:
+        await asyncio.sleep(0)
+        if not aqueue:
+            continue
+        # this could use a lock, if we want thread safety
+        taken = aqueue.copy()
+        aqueue.clear()
+        # action each operation
+        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
+        if lists:
+            try:
+                results = await store.alist(
+                    [p for op in lists.values() for p in op.prefixes]
+                )
+                for fut, op in lists.items():
+                    fut.set_result({k: results.get(k) for k in op.prefixes})
+            except Exception as e:
+                for fut in lists:
+                    fut.set_exception(e)
+        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
+        if puts:
+            try:
+                await store.aput([w for op in puts.values() for w in op.writes])
+                for fut in puts:
+                    fut.set_result(None)
+            except Exception as e:
+                for fut in puts:
+                    fut.set_exception(e)

libs/langgraph/langgraph/store/memory.py

@@ -0,0 +1,25 @@
+from collections import defaultdict
+from typing import List, Optional
+
+from langgraph.store.base import BaseStore, V
+
+
+class MemoryStore(BaseStore):
+    def __init__(self) -> None:
+        self.data: dict[str, dict[str, V]] = defaultdict(dict)
+
+    def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        return {prefix: self.data[prefix] for prefix in prefixes}
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        return self.list(prefixes)
+
+    def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        for namespace, key, value in writes:
+            if value is None:
+                self.data[namespace].pop(key, None)
+            else:
+                self.data[namespace][key] = value
+
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        return self.put(writes)

libs/langgraph/tests/test_pregel.py

@@ -58,6 +58,7 @@ from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
 from langgraph.graph.state import StateGraph
+from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -65,6 +66,7 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
+from langgraph.store.memory import MemoryStore
 from tests.any_str import AnyStr, ExceptionLike
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
@@ -6202,10 +6204,32 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
+        shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+
+    def assert_shared_value(data: State, config: RunnableConfig) -> State:
+        assert "shared" in data
+        if thread_id := config["configurable"].get("thread_id"):
+            if thread_id == "1":
+                # this is the first thread, so should not see a value
+                assert data["shared"] == {}
+                return {"shared": {"1": {"hello": "world"}}}
+            elif thread_id == "2":
+                # this should get value saved by thread 1
+                assert data["shared"] == {"1": {"hello": "world"}}
+            elif thread_id == "3":
+                # this is a different assistant, so should not see previous value
+                assert data["shared"] == {}
+        return {}
+
+    def tool_two_slow(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " slow", **assert_shared_value(data, config)}
+
+    def tool_two_fast(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " fast", **assert_shared_value(data, config)}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
-    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("tool_two_slow", tool_two_slow)
+    tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
@@ -6223,14 +6247,16 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
 
     with SqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
+            store=MemoryStore(),
+            checkpointer=saver,
+            interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
 
         # missing thread_id
         with pytest.raises(ValueError, match="thread_id"):
             tool_two.invoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
             "my_key": "value ⛰️",
@@ -6282,7 +6308,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
         )
 
-        thread2 = {"configurable": {"thread_id": "2"}}
+        thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
             "my_key": "value",
@@ -6322,7 +6348,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config,
         )
 
-        thread3 = {"configurable": {"thread_id": "3"}}
+        thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == {
             "my_key": "value",

libs/langgraph/tests/test_pregel_async.py

@@ -52,6 +52,7 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
+from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -60,6 +61,7 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
+from langgraph.store.memory import MemoryStore
 from tests.any_str import AnyStr, ExceptionLike
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
@@ -4778,10 +4780,33 @@ async def test_start_branch_then() -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
+        shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+        other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+
+    def assert_shared_value(data: State, config: RunnableConfig) -> State:
+        assert "shared" in data
+        if thread_id := config["configurable"].get("thread_id"):
+            if thread_id == "1":
+                # this is the first thread, so should not see a value
+                assert data["shared"] == {}
+                return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}}
+            elif thread_id == "2":
+                # this should get value saved by thread 1
+                assert data["shared"] == {"1": {"hello": "world"}}
+            elif thread_id == "3":
+                # this is a different assistant, so should not see previous value
+                assert data["shared"] == {}
+        return {}
+
+    def tool_two_slow(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " slow", **assert_shared_value(data, config)}
+
+    def tool_two_fast(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " fast", **assert_shared_value(data, config)}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two_slow", lambda s, config: {"my_key": " slow"})
-    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("tool_two_slow", tool_two_slow)
+    tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
@@ -4798,14 +4823,16 @@ async def test_start_branch_then() -> None:
 
     async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
+            store=MemoryStore(),
+            checkpointer=saver,
+            interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
 
         # missing thread_id
         with pytest.raises(ValueError, match="thread_id"):
             await tool_two.ainvoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
             "my_key": "value",
@@ -4865,7 +4892,7 @@ async def test_start_branch_then() -> None:
             ][-1].config,
         )
 
-        thread2 = {"configurable": {"thread_id": "2"}}
+        thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == {
             "my_key": "value",
@@ -4913,7 +4940,7 @@ async def test_start_branch_then() -> None:
             ][-1].config,
         )
 
-        thread3 = {"configurable": {"thread_id": "3"}}
+        thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == {
             "my_key": "value",

libs/langgraph/tests/test_store.py

@@ -0,0 +1,35 @@
+import asyncio
+from typing import Any, Optional
+
+from pytest_mock import MockerFixture
+
+from langgraph.store.base import BaseStore
+from langgraph.store.batch import AsyncBatchedStore
+
+
+async def test_async_batch_store(mocker: MockerFixture) -> None:
+    aget = mocker.stub()
+    alist = mocker.stub()
+
+    class MockStore(BaseStore):
+        async def aget(
+            self, pairs: list[tuple[str, str]]
+        ) -> dict[tuple[str, str], Optional[dict[str, Any]]]:
+            aget(pairs)
+            return {pair: 1 for pair in pairs}
+
+        async def alist(self, prefixes: list[str]) -> dict[str, dict[str, Any]]:
+            alist(prefixes)
+            return {prefix: {prefix: 1} for prefix in prefixes}
+
+    store = AsyncBatchedStore(MockStore())
+
+    # concurrent calls are batched
+    results = await asyncio.gather(
+        store.alist(["a", "b"]),
+        store.alist(["c", "d"]),
+    )
+    assert results == [{"a": {"a": 1}, "b": {"b": 1}}, {"c": {"c": 1}, "d": {"d": 1}}]
+    assert [c.args for c in alist.call_args_list] == [
+        (["a", "b", "c", "d"],),
+    ]