diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index fe47f0d8f..9c7794c46 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/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. diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index f85b33ba3..cf87be337 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/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, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 5810aaf29..cda357607 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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 diff --git a/libs/langgraph/langgraph/kv/__init__.py b/libs/langgraph/langgraph/kv/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/kv/base.py new file mode 100644 index 000000000..570783ec1 --- /dev/null +++ b/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) diff --git a/libs/langgraph/langgraph/kv/memory.py b/libs/langgraph/langgraph/kv/memory.py new file mode 100644 index 000000000..387ea1e7a --- /dev/null +++ b/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) diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py index 0455ed58b..5820383f8 100644 --- a/libs/langgraph/langgraph/managed/base.py +++ b/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() diff --git a/libs/langgraph/langgraph/managed/scoped_value.py b/libs/langgraph/langgraph/managed/scoped_value.py new file mode 100644 index 000000000..adb16e1b4 --- /dev/null +++ b/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)) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d2f0790cc..833453197 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 5922fd14e..994af018e 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f48f86fcd..aaa43be67 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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" diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py new file mode 100644 index 000000000..2e24ee07d --- /dev/null +++ b/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")],), + ] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 138d4dba5..3cc1c4d23 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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",