diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index f85b33ba3..51b6e1437 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_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, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 5810aaf29..8c85b7760 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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 diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py index 0455ed58b..8dbf5c30c 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,7 @@ async def AsyncManagedValuesManager( yield {tasks[task]: task.result() for task in done} else: yield {} + + +ChannelKeyPlaceholder = object() +ChannelTypePlaceholder = object() diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py new file mode 100644 index 000000000..7bb6e23b7 --- /dev/null +++ b/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)) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d2f0790cc..429226214 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 5922fd14e..37f168bf7 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f48f86fcd..672d42049 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 @@ -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" diff --git a/libs/langgraph/langgraph/store/__init__.py b/libs/langgraph/langgraph/store/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/libs/langgraph/langgraph/store/base.py b/libs/langgraph/langgraph/store/base.py new file mode 100644 index 000000000..7f0030f56 --- /dev/null +++ b/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 diff --git a/libs/langgraph/langgraph/store/batch.py b/libs/langgraph/langgraph/store/batch.py new file mode 100644 index 000000000..54eb20d47 --- /dev/null +++ b/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) diff --git a/libs/langgraph/langgraph/store/memory.py b/libs/langgraph/langgraph/store/memory.py new file mode 100644 index 000000000..48fa2884f --- /dev/null +++ b/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) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 138d4dba5..57df91165 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7c5742aa0..70ce65daa 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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", diff --git a/libs/langgraph/tests/test_store.py b/libs/langgraph/tests/test_store.py new file mode 100644 index 000000000..cd53ee407 --- /dev/null +++ b/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"],), + ]