Add ScopedValue

- state shared between threads

libs/langgraph/langgraph/channels/base.py

@@ -33,6 +33,11 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
     # serialize/deserialize methods
 
+    def tap(self) -> Optional[C]:
+        """Return the current checkpoint of the channel, without consuming it.
+        By default, it just calls checkpoint()."""
+        return self.checkpoint()
+
     @abstractmethod
     def checkpoint(self) -> Optional[C]:
         """Return a serializable representation of the channel's current state.

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_KV = "__pregel_kv"
 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_KV,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,

libs/langgraph/langgraph/graph/state.py

@@ -40,7 +40,14 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
-from langgraph.managed.base import ManagedValue, is_managed_value
+from langgraph.kv.base import BaseKV
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    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
@@ -373,6 +380,8 @@ class StateGraph(Graph):
 
     def compile(
         self,
+        *,
+        kv: Optional[BaseKV] = None,
         checkpointer: Optional[BaseCheckpointSaver] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
@@ -442,6 +451,7 @@ class StateGraph(Graph):
             interrupt_after_nodes=interrupt_after,
             auto_validate=False,
             debug=debug,
+            kv=kv,
         )
 
         compiled.attach_node(START, None)
@@ -511,7 +521,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
@@ -684,7 +698,7 @@ def _get_channels(
         return {"__root__": _get_channel(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 +709,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 +750,16 @@ 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
                 return decoration
 
     return None

libs/langgraph/langgraph/kv/base.py

@@ -0,0 +1,113 @@
+import asyncio
+from typing import Any, List, NamedTuple, Optional, Union
+
+V = dict[str, Any]
+
+
+class BaseKV:
+    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], Optional[V]]:
+        # list[(namespace, key)] -> dict[(namespace, key), value | none]
+        raise NotImplementedError
+
+    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 aget(
+        self, pairs: List[tuple[str, str]]
+    ) -> dict[tuple[str, str], Optional[V]]:
+        # list[(namespace, key)] -> dict[(namespace, key), value | 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
+
+
+class GetOp(NamedTuple):
+    pairs: List[tuple[str, str]]
+
+
+class ListOp(NamedTuple):
+    prefixes: List[str]
+
+
+class PutOp(NamedTuple):
+    writes: List[tuple[str, str, Optional[V]]]
+
+
+class KeyValueStore(BaseKV):
+    def __init__(self, kv: BaseKV) -> None:
+        self.kv = kv
+        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
+        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
+
+    def __del__(self) -> None:
+        self.task.cancel()
+
+    async def aget(
+        self, pairs: List[tuple[str, str]]
+    ) -> dict[tuple[str, str], Optional[V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = GetOp(pairs)
+        return await fut
+
+    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[GetOp, ListOp, PutOp]], kv: BaseKV
+) -> 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
+        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
+        if gets:
+            try:
+                results = await kv.aget([p for op in gets.values() for p in op.pairs])
+                for fut, op in gets.items():
+                    fut.set_result({k: results.get(k) for k in op.pairs})
+            except Exception as e:
+                for fut in gets:
+                    fut.set_exception(e)
+        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
+        if lists:
+            try:
+                results = await kv.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 kv.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/kv/memory.py

@@ -0,0 +1,33 @@
+from collections import defaultdict
+from typing import List
+
+from langgraph.kv.base import BaseKV, V
+
+
+class MemoryKV(BaseKV):
+    def __init__(self) -> None:
+        self.data: dict[str, dict[str, V]] = defaultdict(dict)
+
+    def get(self, pairs: List[tuple[str, str]]) -> dict[tuple[str, str], V | None]:
+        return {pair: self.data[pair[0]].get(pair[1]) for pair in pairs}
+
+    async def aget(
+        self, pairs: List[tuple[str, str]]
+    ) -> dict[tuple[str, str], V | None]:
+        return self.get(pairs)
+
+    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, V | None]]) -> 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, V | None]]) -> None:
+        self.put(writes)

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,6 @@ async def AsyncManagedValuesManager(
             yield {tasks[task]: task.result() for task in done}
     else:
         yield {}
+
+
+ChannelKeyPlaceholder = object()

libs/langgraph/langgraph/managed/scoped_value.py

@@ -0,0 +1,99 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterator,
+    Optional,
+    Self,
+    Sequence,
+)
+
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.constants import CONFIG_KEY_KV
+from langgraph.errors import InvalidUpdateError
+from langgraph.kv.base import BaseKV
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ConfiguredManagedValue,
+    WritableManagedValue,
+)
+from langgraph.pregel.types import PregelTaskDescription
+
+V = dict[str, Any]
+
+
+Value = dict[str, V]
+Update = dict[str, Optional[V]]
+
+
+class ScopedValue(WritableManagedValue[Value, Update]):
+    @staticmethod
+    def configure(scope: str) -> ConfiguredManagedValue:
+        return ConfiguredManagedValue(
+            ScopedValue, {"scope": scope, "key": ChannelKeyPlaceholder}
+        )
+
+    @classmethod
+    @contextmanager
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
+        with super().enter(config, **kwargs) as value:
+            if value.kv is not None:
+                saved = value.kv.list([value.ns])
+                value.value = saved[value.ns]
+            yield value
+
+    @classmethod
+    @asynccontextmanager
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
+        async with super().aenter(config, **kwargs) as value:
+            if value.kv is not None:
+                saved = await value.kv.alist([value.ns])
+                value.value = saved[value.ns]
+            yield value
+
+    def __init__(self, config: RunnableConfig, *, scope: str, key: str) -> None:
+        self.scope = scope
+        self.config = config
+        self.value: Value = {}
+        self.kv: BaseKV = config["configurable"].get(CONFIG_KEY_KV)
+        if self.kv 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, task: PregelTaskDescription) -> 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.kv is None:
+            self._process_update(values)
+        else:
+            return self.kv.put(self._process_update(values))
+
+    async def aupdate(self, writes: Sequence[Update]) -> None:
+        if self.kv is None:
+            self._process_update(writes)
+        else:
+            return await self.kv.aput(self._process_update(writes))

libs/langgraph/langgraph/pregel/__init__.py

@@ -73,6 +73,7 @@ from langgraph.constants import (
     Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
+from langgraph.kv.base import BaseKV
 from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValuesManager,
@@ -223,6 +224,9 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
+    kv: Optional[BaseKV] = None
+    """Key-value store to use. Defaults to None."""
+
     retry_policy: Optional[RetryPolicy] = None
     """Retry policy to use when running tasks. Set to None to disable."""
 
@@ -644,7 +648,7 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             )
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
@@ -788,7 +792,7 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             )
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)

libs/langgraph/langgraph/pregel/algo.py

@@ -145,7 +145,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 +161,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 +178,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 +218,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(

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
@@ -40,6 +41,7 @@ from langgraph.checkpoint.base import (
     empty_checkpoint,
 )
 from langgraph.constants import (
+    CONFIG_KEY_KV,
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
     ERROR,
@@ -52,6 +54,7 @@ from langgraph.managed.base import (
     AsyncManagedValuesManager,
     ManagedValueMapping,
     ManagedValuesManager,
+    WritableManagedValue,
 )
 from langgraph.pregel.algo import (
     PregelTaskWrites,
@@ -182,12 +185,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 +330,13 @@ class PregelLoop:
                 manager=None,
             )
             # apply input writes
-            apply_writes(
+            mv_writes = apply_writes(
                 self.checkpoint,
                 self.channels,
                 discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
                 self.checkpointer_get_next_version,
             )
+            assert not mv_writes
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": self.input})
         else:
@@ -395,6 +402,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]],
@@ -438,6 +448,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 +474,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_KV: self.graph.kv}),
+            )
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
@@ -515,6 +531,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 +561,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_KV: self.graph.kv}),
+            )
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"

libs/langgraph/tests/test_kv.py

@@ -0,0 +1,32 @@
+import asyncio
+from typing import Any, List
+
+from pytest_mock import MockerFixture
+
+from langgraph.kv.base import BaseKV, KeyValueStore
+
+
+async def test_kv_queue(mocker: MockerFixture) -> None:
+    aget = mocker.stub()
+
+    class MockKV(BaseKV):
+        async def aget(
+            self, pairs: List[tuple[str, str]]
+        ) -> dict[tuple[str, str], dict[str, Any] | None]:
+            aget(pairs)
+            return {pair: {0: pair[0], 1: pair[1]} for pair in pairs}
+
+    store = KeyValueStore(MockKV())
+
+    # concurrent calls are batched
+    results = await asyncio.gather(
+        store.aget([("a", "b")]),
+        store.aget([("c", "d")]),
+    )
+    assert results == [
+        {("a", "b"): {0: "a", 1: "b"}},
+        {("c", "d"): {0: "c", 1: "d"}},
+    ]
+    assert [c.args for c in aget.call_args_list] == [
+        ([("a", "b"), ("c", "d")],),
+    ]

libs/langgraph/tests/test_pregel.py

@@ -58,6 +58,8 @@ 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.kv.memory import MemoryKV
+from langgraph.managed.scoped_value import ScopedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -165,6 +167,7 @@ def test_graph_validation() -> None:
 
     class State(TypedDict):
         hello: str
+        shared_things: Annotated[dict[str, dict[str, Any]], ScopedValue("assistant_id")]
 
     def node_a(state: State) -> State:
         # typo
@@ -6202,10 +6205,34 @@ 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]], ScopedValue.configure("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 +6250,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"]
+            kv=MemoryKV(),
+            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 +6311,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 +6351,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",