diff --git a/libs/cli/langgraph_cli/docker.py b/libs/cli/langgraph_cli/docker.py index 79e0dd25a..d2cffc26b 100644 --- a/libs/cli/langgraph_cli/docker.py +++ b/libs/cli/langgraph_cli/docker.py @@ -12,6 +12,15 @@ DEFAULT_POSTGRES_URI = ( "postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable" ) +REDIS = """ + langgraph-redis: + image: redis:6 + healthcheck: + test: redis-cli ping + interval: 5s + timeout: 1s + retries: 5 +""" DB = """ langgraph-postgres: @@ -166,18 +175,22 @@ def compose( compose_str = f"""{volumes}services: {db} +{REDIS} {debugger_compose(port=debugger_port, base_url=debugger_base_url)} langgraph-api: ports: - - "{port}:8000\"""" + - "{port}:8000\" + depends_on: + langgraph-redis: + condition: service_healthy""" if include_db: compose_str += """ - depends_on: langgraph-postgres: condition: service_healthy""" compose_str += f""" environment: POSTGRES_URI: {postgres_uri} + REDIS_URI: redis://langgraph-redis:6379 """ if capabilities.healthcheck_start_interval: compose_str += """ healthcheck: diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py deleted file mode 100644 index 9a7c8d87e..000000000 --- a/libs/langgraph/langgraph/channels/manager.py +++ /dev/null @@ -1,39 +0,0 @@ -from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager -from typing import AsyncGenerator, Generator, Mapping - -from langchain_core.runnables import RunnableConfig - -from langgraph.channels.base import BaseChannel -from langgraph.checkpoint.base import Checkpoint - - -@contextmanager -def ChannelsManager( - channels: Mapping[str, BaseChannel], - checkpoint: Checkpoint, - config: RunnableConfig, -) -> Generator[Mapping[str, BaseChannel], None, None]: - """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" - with ExitStack() as stack: - yield { - k: stack.enter_context( - v.from_checkpoint(checkpoint["channel_values"].get(k), config) - ) - for k, v in channels.items() - } - - -@asynccontextmanager -async def AsyncChannelsManager( - channels: Mapping[str, BaseChannel], - checkpoint: Checkpoint, - config: RunnableConfig, -) -> AsyncGenerator[Mapping[str, BaseChannel], None]: - """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" - async with AsyncExitStack() as stack: - yield { - k: await stack.enter_async_context( - v.afrom_checkpoint(checkpoint["channel_values"].get(k), config) - ) - for k, v in channels.items() - } diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index d0c1ebff7..603de9395 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -6,6 +6,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__" @@ -18,6 +19,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 f998b568a..259d51b62 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -44,10 +44,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, + ManagedValueSpec, + 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__) @@ -125,8 +133,8 @@ class StateGraph(Graph): nodes: dict[str, StateNodeSpec] channels: dict[str, BaseChannel] - managed: dict[str, Type[ManagedValue]] - schemas: dict[Type[Any], dict[str, Union[BaseChannel, Type[ManagedValue]]]] + managed: dict[str, ManagedValueSpec] + schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]] def __init__( self, @@ -382,6 +390,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, @@ -440,7 +450,11 @@ class StateGraph(Graph): builder=self, config_type=self.config_schema, nodes={}, - channels={**self.channels, START: EphemeralValue(self.input)}, + channels={ + **self.channels, + **self.managed, + START: EphemeralValue(self.input), + }, input_channels=START, stream_mode="updates", output_channels=output_channels, @@ -450,6 +464,7 @@ class StateGraph(Graph): interrupt_after_nodes=interrupt_after, auto_validate=False, debug=debug, + store=store, ) compiled.attach_node(START, None) @@ -494,7 +509,7 @@ class CompiledStateGraph(CompiledGraph): **{ k: (self.channels[k].UpdateType, None) for k in self.builder.schemas[self.builder.input] - if k in self.channels + if isinstance(self.channels[k], BaseChannel) and not isinstance(self.channels[k], Context) }, ) @@ -519,7 +534,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 @@ -565,10 +584,7 @@ class CompiledStateGraph(CompiledGraph): ) else: input_schema = node.input if node else self.builder.schema - input_values = { - k: v if is_managed_value(v) else k - for k, v in self.builder.schemas[input_schema].items() - } + input_values = {k: k for k in self.builder.schemas[input_schema]} is_single_input = len(input_values) == 1 and "__root__" in input_values self.channels[key] = EphemeralValue(Any, guard=False) @@ -687,12 +703,12 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: def _get_channels( schema: Type[dict], -) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]: +) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]: 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__" } @@ -703,9 +719,9 @@ def _get_channels( def _get_channel( - annotation: Any, *, allow_managed: bool = True -) -> Union[BaseChannel, Type[ManagedValue]]: - if manager := _is_field_managed_value(annotation): + name: str, annotation: Any, *, allow_managed: bool = True +) -> Union[BaseChannel, ManagedValueSpec]: + if manager := _is_field_managed_value(name, annotation): if allow_managed: return manager else: @@ -744,12 +760,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[ManagedValueSpec]: 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..bebca1be2 100644 --- a/libs/langgraph/langgraph/managed/base.py +++ b/libs/langgraph/langgraph/managed/base.py @@ -1,13 +1,13 @@ -import asyncio from abc import ABC, abstractmethod -from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager +from contextlib import asynccontextmanager, contextmanager from inspect import isclass from typing import ( Any, - AsyncGenerator, - Generator, + AsyncIterator, Generic, + Iterator, NamedTuple, + Sequence, Type, TypeVar, Union, @@ -17,6 +17,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 +26,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 +40,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 +57,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,46 +83,23 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]: ) -@contextmanager -def ManagedValuesManager( - values: dict[str, ManagedValueSpec], - config: RunnableConfig, -) -> Generator[ManagedValueMapping, None, None]: - if values: - with ExitStack() as stack: - yield { - key: stack.enter_context( - value.cls.enter(config, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.enter(config) - ) - for key, value in values.items() - } - else: - yield {} +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) + ) -@asynccontextmanager -async def AsyncManagedValuesManager( - values: dict[str, ManagedValueSpec], - config: RunnableConfig, -) -> AsyncGenerator[ManagedValueMapping, None]: - if values: - async with AsyncExitStack() as stack: - # create enter tasks with reference to spec - tasks = { - asyncio.create_task( - stack.enter_async_context( - value.cls.aenter(config, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.aenter(config) - ) - ): key - for key, value in values.items() - } - # wait for all enter tasks - done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED) - # build mapping from spec to result - yield {tasks[task]: task.result() for task in done} - else: - yield {} +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) + ) + + +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 5c54d5b7e..ed3ec185a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -52,14 +52,8 @@ from langgraph.channels.base import ( BaseChannel, ) from langgraph.channels.context import Context -from langgraph.channels.last_value import LastValue -from langgraph.channels.manager import ( - AsyncChannelsManager, - ChannelsManager, -) from langgraph.checkpoint.base import ( BaseCheckpointSaver, - CheckpointTuple, copy_checkpoint, create_checkpoint, empty_checkpoint, @@ -76,12 +70,7 @@ from langgraph.constants import ( Interrupt, ) from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError -from langgraph.managed.base import ( - AsyncManagedValuesManager, - ManagedValuesManager, - ManagedValueSpec, - is_managed_value, -) +from langgraph.managed.base import ManagedValueSpec from langgraph.pregel.algo import ( apply_writes, local_read, @@ -100,6 +89,7 @@ from langgraph.pregel.io import ( read_channels, ) from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop +from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry from langgraph.pregel.types import ( @@ -111,6 +101,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], @@ -273,7 +264,9 @@ class Pregel( ): nodes: Mapping[str, PregelNode] - channels: Mapping[str, BaseChannel] = Field(default_factory=dict) + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field( + default_factory=dict + ) auto_validate: bool = True @@ -300,6 +293,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.""" @@ -420,19 +416,12 @@ class Pregel( @property def stream_channels_asis(self) -> Union[str, Sequence[str]]: return self.stream_channels or [ - k for k in self.channels if not isinstance(self.channels[k], Context) + k + for k in self.channels + if isinstance(self.channels[k], BaseChannel) + and not isinstance(self.channels[k], Context) ] - @property - def managed_values_dict(self) -> dict[str, ManagedValueSpec]: - return { - k: v - for node in self.nodes.values() - if isinstance(node.channels, dict) - for k, v in node.channels.items() - if is_managed_value(v) - } - @property def subgraphs(self) -> Iterator[Pregel]: for node in self.nodes.values(): @@ -478,15 +467,11 @@ class Pregel( graph = checkpoint_ns_to_graph[saved_checkpoint_ns] with ChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in graph.channels.items() - }, - saved.checkpoint, - saved.config, - ) as channels, ManagedValuesManager( - graph.managed_values_dict, ensure_config(saved.config) - ) as managed: + graph.channels, saved.checkpoint, saved.config, skip_context=True + ) as ( + channels, + managed, + ): next_tasks = prepare_next_tasks( saved.checkpoint, graph.nodes, @@ -547,6 +532,7 @@ class Pregel( existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get( saved_checkpoint_ns ) + # keep only most recent checkpoint_id if ( existing_checkpoint_id is None @@ -559,15 +545,8 @@ class Pregel( graph = checkpoint_ns_to_graph[saved_checkpoint_ns] async with AsyncChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in graph.channels.items() - }, - saved.checkpoint, - saved.config, - ) as channels, AsyncManagedValuesManager( - graph.managed_values_dict, ensure_config(saved.config) - ) as managed: + graph.channels, saved.checkpoint, saved.config, skip_context=True + ) as (channels, managed): next_tasks = prepare_next_tasks( saved.checkpoint, graph.nodes, @@ -742,11 +721,10 @@ class Pregel( if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") # update channels - with ChannelsManager( - self.channels, checkpoint, config - ) as channels, ManagedValuesManager( - self.managed_values_dict, ensure_config(config) - ) as managed: + with ChannelsManager(self.channels, checkpoint, config) as ( + channels, + managed, + ): # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() if not writers: @@ -777,9 +755,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( @@ -886,11 +864,10 @@ class Pregel( if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") # update channels, acting as the chosen node - async with AsyncChannelsManager( - self.channels, checkpoint, config - ) as channels, AsyncManagedValuesManager( - self.managed_values_dict, ensure_config(config) - ) as managed: + async with AsyncChannelsManager(self.channels, checkpoint, config) as ( + channels, + managed, + ): # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() if not writers: @@ -921,9 +898,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( @@ -1127,7 +1104,13 @@ class Pregel( ) with SyncPregelLoop( - input, config=config, checkpointer=checkpointer, graph=self + input, + config=config, + store=self.store, + checkpointer=checkpointer, + nodes=self.nodes, + specs=self.channels, + output_keys=output_keys, ) as loop: # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates @@ -1135,7 +1118,8 @@ class Pregel( # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick( - output_keys=output_keys, + input_keys=self.input_channels, + stream_keys=self.stream_channels_asis, interrupt_before=interrupt_before, interrupt_after=interrupt_after, manager=run_manager, @@ -1255,8 +1239,8 @@ class Pregel( "without hitting a stop condition. You can increase the " "limit by setting the `recursion_limit` config key." ) - # set final channel values as run output - run_manager.on_chain_end(read_channels(loop.channels, output_keys)) + # set final channel values as run output + run_manager.on_chain_end(loop.output) except BaseException as e: run_manager.on_chain_error(e) raise @@ -1379,7 +1363,13 @@ class Pregel( debug=debug, ) async with AsyncPregelLoop( - input, config=config, checkpointer=checkpointer, graph=self + input, + config=config, + store=self.store, + checkpointer=checkpointer, + nodes=self.nodes, + specs=self.channels, + output_keys=output_keys, ) as loop: aioloop = asyncio.get_event_loop() # Similarly to Bulk Synchronous Parallel / Pregel model @@ -1388,7 +1378,8 @@ class Pregel( # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick( - output_keys=output_keys, + input_keys=self.input_channels, + stream_keys=self.stream_channels_asis, interrupt_before=interrupt_before, interrupt_after=interrupt_after, manager=run_manager, @@ -1511,13 +1502,9 @@ class Pregel( "without hitting a stop condition. You can increase the " "limit by setting the `recursion_limit` config key." ) - - # set final channel values as run output - await run_manager.on_chain_end( - read_channels(loop.channels, output_keys) - ) + # set final channel values as run output + await run_manager.on_chain_end(loop.output) except BaseException as e: - # TODO use on_chain_end if exc is GraphInterrupt await asyncio.shield(run_manager.on_chain_error(e)) raise diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 453bd6a76..846afd9e0 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -25,7 +25,6 @@ from langchain_core.runnables.config import ( from langgraph.channels.base import BaseChannel from langgraph.channels.context import Context -from langgraph.channels.manager import ChannelsManager from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -46,9 +45,10 @@ from langgraph.constants import ( Send, ) from langgraph.errors import EmptyChannelError, InvalidUpdateError -from langgraph.managed.base import ManagedValueMapping, is_managed_value +from langgraph.managed.base import ManagedValueMapping from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger +from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All, PregelExecutableTask, PregelTask @@ -105,11 +105,10 @@ def local_read( if fresh: new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1) context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)} - with ChannelsManager( - {k: v for k, v in channels.items() if k not in context_channels}, - new_checkpoint, - config, - ) as channels: + with ChannelsManager(channels, new_checkpoint, config, skip_context=True) as ( + channels, + _, + ): all_channels = {**channels, **context_channels} apply_writes(new_checkpoint, all_channels, [task], None) return read_channels(all_channels, select) @@ -121,6 +120,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 +131,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 +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( @@ -329,7 +336,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, @@ -422,7 +433,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, @@ -471,16 +486,10 @@ def _proc_input( chan, catch=chan not in proc.triggers, ) + if chan in channels + else managed[k](step) for k, chan in proc.channels.items() - if isinstance(chan, str) } - - managed_values = {} - for key, chan in proc.channels.items(): - if is_managed_value(chan): - managed_values[key] = managed[key](step) - - val.update(managed_values) except EmptyChannelError: return elif isinstance(proc.channels, list): diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f48f86fcd..aa256b2f7 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -4,7 +4,6 @@ from collections import deque from contextlib import AsyncExitStack, ExitStack from types import TracebackType from typing import ( - TYPE_CHECKING, Any, AsyncContextManager, Callable, @@ -18,6 +17,7 @@ from typing import ( Type, TypeVar, Union, + cast, ) from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager @@ -25,10 +25,6 @@ from langchain_core.runnables import RunnableConfig from typing_extensions import Self from langgraph.channels.base import BaseChannel -from langgraph.channels.manager import ( - AsyncChannelsManager, - ChannelsManager, -) from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -49,9 +45,9 @@ from langgraph.constants import ( ) from langgraph.errors import EmptyInputError, GraphInterrupt from langgraph.managed.base import ( - AsyncManagedValuesManager, ManagedValueMapping, - ManagedValuesManager, + ManagedValueSpec, + WritableManagedValue, ) from langgraph.pregel.algo import ( PregelTaskWrites, @@ -66,13 +62,19 @@ from langgraph.pregel.executor import ( BackgroundExecutor, Submit, ) -from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single +from langgraph.pregel.io import ( + map_input, + map_output_updates, + map_output_values, + read_channels, + single, +) +from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager +from langgraph.pregel.read import PregelNode from langgraph.pregel.types import PregelExecutableTask from langgraph.pregel.utils import get_new_channel_versions - -if TYPE_CHECKING: - from langgraph.pregel import Pregel - +from langgraph.store.base import BaseStore +from langgraph.store.batch import AsyncBatchedStore V = TypeVar("V") INPUT_DONE = object() @@ -83,7 +85,13 @@ EMPTY_SEQ = () class PregelLoop: input: Optional[Any] config: RunnableConfig + store: Optional[BaseStore] checkpointer: Optional[BaseCheckpointSaver] + nodes: Mapping[str, PregelNode] + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] + output_keys: Union[str, Sequence[str]] + is_nested: bool + checkpointer_get_next_version: Callable[[Optional[V]], V] checkpointer_put_writes: Optional[ Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any] @@ -99,8 +107,6 @@ class PregelLoop: Any, ] ] - graph: "Pregel" - submit: Submit channels: Mapping[str, BaseChannel] managed: ManagedValueMapping @@ -118,7 +124,7 @@ class PregelLoop: ] tasks: Sequence[PregelExecutableTask] stream: deque[Tuple[str, Any]] - is_nested: bool + output: Union[None, dict[str, Any], Any] = None # public @@ -127,16 +133,20 @@ class PregelLoop: input: Optional[Any], *, config: RunnableConfig, + store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], - graph: "Pregel", + nodes: Mapping[str, PregelNode], + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + output_keys: Union[str, Sequence[str]], ) -> None: self.stream = deque() self.input = input self.config = config + self.store = store self.checkpointer = checkpointer - self.graph = graph - # TODO if managed values no longer needs graph we can replace with - # managed_specs, channel_specs + self.nodes = nodes + self.specs = specs + self.output_keys = output_keys self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: @@ -166,7 +176,8 @@ class PregelLoop: def tick( self, *, - output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, + input_keys: Union[str, Sequence[str]], + stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, interrupt_after: Sequence[str] = EMPTY_SEQ, interrupt_before: Sequence[str] = EMPTY_SEQ, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, @@ -178,20 +189,23 @@ class PregelLoop: raise RuntimeError("Cannot tick when status is no longer 'pending'") if self.input not in (INPUT_DONE, INPUT_RESUMING): - self._first() + self._first(input_keys=input_keys) 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) - for v in map_output_values(output_keys, writes, self.channels) + for v in map_output_values(self.output_keys, writes, self.channels) ) # clear pending writes self.checkpoint_pending_writes.clear() @@ -199,11 +213,7 @@ class PregelLoop: self._put_checkpoint( { "source": "loop", - "writes": single( - map_output_updates(output_keys, self.tasks) - if self.graph.stream_mode == "updates" - else map_output_values(output_keys, writes, self.channels) - ), + "writes": single(map_output_updates(self.output_keys, self.tasks)), } ) # after execution, check if we should interrupt @@ -227,7 +237,7 @@ class PregelLoop: # prepare next tasks self.tasks = prepare_next_tasks( self.checkpoint, - self.graph.nodes, + self.nodes, self.channels, self.managed, self.config, @@ -246,7 +256,7 @@ class PregelLoop: self.step - 1, # printing checkpoint for previous step self.checkpoint_config, self.channels, - self.graph.stream_channels_asis, + stream_keys, self.checkpoint_metadata, self.checkpoint, self.tasks, @@ -270,7 +280,8 @@ class PregelLoop: # if all tasks have finished, re-tick if all(task.writes for task in self.tasks): return self.tick( - output_keys=output_keys, + input_keys=input_keys, + stream_keys=stream_keys, interrupt_after=interrupt_after, interrupt_before=interrupt_before, manager=manager, @@ -294,7 +305,7 @@ class PregelLoop: # private - def _first(self) -> None: + def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None: # resuming from previous checkpoint requires # - finding a previous checkpoint # - receiving None input (outer graph) or RESUMING flag (subgraph) @@ -311,11 +322,11 @@ class PregelLoop: version = self.checkpoint["channel_versions"][k] self.checkpoint["versions_seen"][INTERRUPT][k] = version # map inputs to channel updates - elif input_writes := deque(map_input(self.graph.input_channels, self.input)): + elif input_writes := deque(map_input(input_keys, self.input)): # discard any unfinished tasks from previous checkpoint discard_tasks = prepare_next_tasks( self.checkpoint, - self.graph.nodes, + self.nodes, self.channels, self.managed, self.config, @@ -324,16 +335,16 @@ 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: - raise EmptyInputError(f"Received no input for {self.graph.input_channels}") + raise EmptyInputError(f"Received no input for {input_keys}") # done with input self.input = INPUT_RESUMING if is_resuming else INPUT_DONE @@ -395,13 +406,21 @@ 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]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - if isinstance(exc_value, GraphInterrupt) and not self.is_nested: + suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested + if suppress or exc_type is None: + # save final output + self.output = read_channels(self.channels, self.output_keys) + if suppress: + # suppress interrupt return True @@ -411,10 +430,21 @@ class SyncPregelLoop(PregelLoop, ContextManager): input: Optional[Any], *, config: RunnableConfig, + store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], - graph: "Pregel", + nodes: Mapping[str, PregelNode], + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, ) -> None: - super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) + super().__init__( + input, + config=config, + checkpointer=checkpointer, + store=store, + nodes=nodes, + specs=specs, + output_keys=output_keys, + ) self.stack = ExitStack() if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version @@ -438,6 +468,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: @@ -457,11 +490,8 @@ class SyncPregelLoop(PregelLoop, ContextManager): self.checkpoint_pending_writes = saved.pending_writes or [] self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) - self.channels = self.stack.enter_context( - ChannelsManager(self.graph.channels, self.checkpoint, self.config) - ) - self.managed = self.stack.enter_context( - ManagedValuesManager(self.graph.managed_values_dict, self.config) + self.channels, self.managed = self.stack.enter_context( + ChannelsManager(self.specs, self.checkpoint, self.config, self.store) ) self.stack.push(self._suppress_interrupt) self.status = "pending" @@ -478,7 +508,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): traceback: Optional[TracebackType], ) -> Optional[bool]: # unwind stack - del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) @@ -488,10 +517,22 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): input: Optional[Any], *, config: RunnableConfig, + store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], - graph: "Pregel", + nodes: Mapping[str, PregelNode], + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, ) -> None: - super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) + super().__init__( + input, + config=config, + checkpointer=checkpointer, + store=store, + nodes=nodes, + specs=specs, + output_keys=output_keys, + ) + self.store = AsyncBatchedStore(self.store) if self.store else None self.stack = AsyncExitStack() if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version @@ -515,6 +556,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: @@ -536,11 +582,8 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): self.checkpoint_pending_writes = saved.pending_writes or [] self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor()) - self.channels = await self.stack.enter_async_context( - AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config) - ) - self.managed = await self.stack.enter_async_context( - AsyncManagedValuesManager(self.graph.managed_values_dict, self.config) + self.channels, self.managed = await self.stack.enter_async_context( + AsyncChannelsManager(self.specs, self.checkpoint, self.config, self.store) ) self.stack.push(self._suppress_interrupt) self.status = "pending" @@ -558,7 +601,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): traceback: Optional[TracebackType], ) -> Optional[bool]: # unwind stack - del self.graph return await asyncio.shield( self.stack.__aexit__(exc_type, exc_value, traceback) ) diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py new file mode 100644 index 000000000..437019113 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/manager.py @@ -0,0 +1,104 @@ +import asyncio +from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager +from typing import AsyncIterator, Iterator, Mapping, Optional, Union + +from langchain_core.runnables import RunnableConfig, patch_config + +from langgraph.channels.base import BaseChannel +from langgraph.channels.context import Context +from langgraph.channels.last_value import LastValue +from langgraph.checkpoint.base import Checkpoint +from langgraph.constants import CONFIG_KEY_STORE +from langgraph.managed.base import ( + ConfiguredManagedValue, + ManagedValueMapping, + ManagedValueSpec, +) +from langgraph.store.base import BaseStore + + +@contextmanager +def ChannelsManager( + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + checkpoint: Checkpoint, + config: RunnableConfig, + store: Optional[BaseStore] = None, + *, + skip_context: bool = False, +) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: + """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" + config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store}) + channel_specs: Mapping[str, BaseChannel] = {} + managed_specs: Mapping[str, ManagedValueSpec] = {} + for k, v in specs.items(): + if skip_context and isinstance(v, Context): + channel_specs[k] = LastValue(None) + elif isinstance(v, BaseChannel): + channel_specs[k] = v + else: + managed_specs[k] = v + with ExitStack() as stack: + yield ( + { + k: stack.enter_context( + v.from_checkpoint(checkpoint["channel_values"].get(k), config) + ) + for k, v in channel_specs.items() + }, + { + key: stack.enter_context( + value.cls.enter(config_for_managed, **value.kwargs) + if isinstance(value, ConfiguredManagedValue) + else value.enter(config_for_managed) + ) + for key, value in managed_specs.items() + }, + ) + + +@asynccontextmanager +async def AsyncChannelsManager( + specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + checkpoint: Checkpoint, + config: RunnableConfig, + store: Optional[BaseStore] = None, + *, + skip_context: bool = False, +) -> AsyncIterator[Mapping[str, BaseChannel]]: + """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" + config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store}) + channel_specs: Mapping[str, BaseChannel] = {} + managed_specs: Mapping[str, ManagedValueSpec] = {} + for k, v in specs.items(): + if skip_context and isinstance(v, Context): + channel_specs[k] = LastValue(None) + elif isinstance(v, BaseChannel): + channel_specs[k] = v + else: + managed_specs[k] = v + async with AsyncExitStack() as stack: + # managed: create enter tasks with reference to spec, await them + if tasks := { + asyncio.create_task( + stack.enter_async_context( + value.cls.aenter(config_for_managed, **value.kwargs) + if isinstance(value, ConfiguredManagedValue) + else value.aenter(config_for_managed) + ) + ): key + for key, value in managed_specs.items() + }: + done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED) + else: + done = set() + yield ( + # channels: enter each channel with checkpoint + { + k: await stack.enter_async_context( + v.afrom_checkpoint(checkpoint["channel_values"].get(k), config) + ) + for k, v in channel_specs.items() + }, + # managed: build mapping from spec to result + {tasks[task]: task.result() for task in done}, + ) diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 6e492e2ac..b5c971e4a 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -15,7 +15,6 @@ from langchain_core.runnables.config import merge_configs from langchain_core.runnables.utils import ConfigurableFieldSpec from langgraph.constants import CONFIG_KEY_READ -from langgraph.managed.base import ManagedValueSpec from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.write import ChannelWrite from langgraph.utils import RunnableCallable @@ -101,7 +100,7 @@ DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() class PregelNode(RunnableBindingBase): - channels: Union[list[str], Mapping[str, Union[str, ManagedValueSpec]]] + channels: Union[list[str], Mapping[str, str]] triggers: list[str] = Field(default_factory=list) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 88940b584..486584809 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -4,6 +4,7 @@ import random import time from typing import Optional +from langgraph.errors import GraphInterrupt from langgraph.pregel.types import PregelExecutableTask, RetryPolicy logger = logging.getLogger(__name__) @@ -25,6 +26,9 @@ def run_with_retry( task.proc.invoke(task.input, task.config) # if successful, end break + except GraphInterrupt: + # if interrupted, end + raise except Exception as exc: if retry_policy is None: raise @@ -75,6 +79,9 @@ async def arun_with_retry( await task.proc.ainvoke(task.input, task.config) # if successful, end break + except GraphInterrupt: + # if interrupted, end + raise except Exception as exc: if retry_policy is None: raise 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/pyproject.toml b/libs/langgraph/pyproject.toml index 4dff54572..1df951ee4 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.6" +version = "0.2.9" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" @@ -65,7 +65,7 @@ omit = ["tests/*"] [tool.pytest-watcher] now = true delay = 0.1 -runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"] +runner_args = ["--ff", "-v", "-x", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"] patterns = ["*.py"] [build-system] diff --git a/libs/langgraph/tests/fake_tracer.py b/libs/langgraph/tests/fake_tracer.py new file mode 100644 index 000000000..28ecc88db --- /dev/null +++ b/libs/langgraph/tests/fake_tracer.py @@ -0,0 +1,91 @@ +from typing import Any, Optional +from uuid import UUID + +from langchain_core.messages.base import BaseMessage +from langchain_core.outputs.chat_generation import ChatGeneration +from langchain_core.outputs.llm_result import LLMResult +from langchain_core.tracers import BaseTracer, Run + + +class FakeTracer(BaseTracer): + """Fake tracer that records LangChain execution. + It replaces run ids with deterministic UUIDs for snapshotting.""" + + def __init__(self) -> None: + """Initialize the tracer.""" + super().__init__() + self.runs: list[Run] = [] + self.uuids_map: dict[UUID, UUID] = {} + self.uuids_generator = ( + UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000) + ) + + def _replace_uuid(self, uuid: UUID) -> UUID: + if uuid not in self.uuids_map: + self.uuids_map[uuid] = next(self.uuids_generator) + return self.uuids_map[uuid] + + def _replace_message_id(self, maybe_message: Any) -> Any: + if isinstance(maybe_message, BaseMessage): + maybe_message.id = str(next(self.uuids_generator)) + if isinstance(maybe_message, ChatGeneration): + maybe_message.message.id = str(next(self.uuids_generator)) + if isinstance(maybe_message, LLMResult): + for i, gen_list in enumerate(maybe_message.generations): + for j, gen in enumerate(gen_list): + maybe_message.generations[i][j] = self._replace_message_id(gen) + if isinstance(maybe_message, dict): + for k, v in maybe_message.items(): + maybe_message[k] = self._replace_message_id(v) + if isinstance(maybe_message, list): + for i, v in enumerate(maybe_message): + maybe_message[i] = self._replace_message_id(v) + + return maybe_message + + def _copy_run(self, run: Run) -> Run: + if run.dotted_order: + levels = run.dotted_order.split(".") + processed_levels = [] + for level in levels: + timestamp, run_id = level.split("Z") + new_run_id = self._replace_uuid(UUID(run_id)) + processed_level = f"{timestamp}Z{new_run_id}" + processed_levels.append(processed_level) + new_dotted_order = ".".join(processed_levels) + else: + new_dotted_order = None + return run.copy( + update={ + "id": self._replace_uuid(run.id), + "parent_run_id": ( + self.uuids_map[run.parent_run_id] if run.parent_run_id else None + ), + "child_runs": [self._copy_run(child) for child in run.child_runs], + "trace_id": self._replace_uuid(run.trace_id) if run.trace_id else None, + "dotted_order": new_dotted_order, + "inputs": self._replace_message_id(run.inputs), + "outputs": self._replace_message_id(run.outputs), + } + ) + + def _persist_run(self, run: Run) -> None: + """Persist a run.""" + + self.runs.append(self._copy_run(run)) + + def flattened_runs(self) -> list[Run]: + q = [] + self.runs + result = [] + while q: + parent = q.pop() + result.append(parent) + if parent.child_runs: + q.extend(parent.child_runs) + return result + + @property + def run_ids(self) -> list[Optional[UUID]]: + runs = self.flattened_runs() + uuids_map = {v: k for k, v in self.uuids_map.items()} + return [uuids_map.get(r.id) for r in runs] diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index cdf179c24..2102261ae 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -1,7 +1,6 @@ -from langgraph.channels.manager import ChannelsManager from langgraph.checkpoint.base import empty_checkpoint -from langgraph.managed.base import ManagedValuesManager from langgraph.pregel.algo import prepare_next_tasks +from langgraph.pregel.manager import ChannelsManager def test_prepare_next_tasks() -> None: @@ -9,9 +8,7 @@ def test_prepare_next_tasks() -> None: processes = {} checkpoint = empty_checkpoint() - with ManagedValuesManager({}, config) as managed, ChannelsManager( - {}, checkpoint, config - ) as channels: + with ChannelsManager({}, checkpoint, config) as (channels, managed): assert ( prepare_next_tasks( checkpoint, processes, channels, managed, config, 0, for_execution=False diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 6058ecf1d..69842eb9a 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, ) @@ -70,7 +71,9 @@ from langgraph.pregel import ( ) 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.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, @@ -660,7 +663,7 @@ def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 6, "writes": 5}, + metadata={"source": "loop", "step": 6, "writes": {"two": 5}}, created_at=AnyStr(), parent_config=history[1].config, ), @@ -675,7 +678,7 @@ def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 5, "writes": None}, + metadata={"source": "loop", "step": 5, "writes": {"one": None}}, created_at=AnyStr(), parent_config=history[2].config, ), @@ -705,7 +708,7 @@ def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 3, "writes": None}, + metadata={"source": "loop", "step": 3, "writes": {"one": None}}, created_at=AnyStr(), parent_config=history[4].config, ), @@ -735,7 +738,7 @@ def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 1, "writes": 4}, + metadata={"source": "loop", "step": 1, "writes": {"two": 4}}, created_at=AnyStr(), parent_config=history[6].config, ), @@ -750,7 +753,7 @@ def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 0, "writes": None}, + metadata={"source": "loop", "step": 0, "writes": {"one": None}}, created_at=AnyStr(), parent_config=history[7].config, ), @@ -1998,12 +2001,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "step": 0, "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } }, }, }, @@ -2209,12 +2214,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "step": 0, "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } } }, }, @@ -2414,12 +2421,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "step": 0, "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } } }, }, @@ -6142,20 +6151,34 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: my_key: Annotated[str, operator.add] market: str + tool_two_node_count = 0 + def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 if s["market"] == "DE": raise NodeInterrupt("Just because...") return {"my_key": " all good"} tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two", tool_two_node) + tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) tool_two_graph.add_edge(START, "tool_two") tool_two = tool_two_graph.compile() - assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { + tracer = FakeTracer() + assert tool_two.invoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { "my_key": "value", "market": "DE", } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { "my_key": "value all good", "market": "US", @@ -6207,10 +6230,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 ) @@ -6228,14 +6273,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 ⛰️", @@ -6287,7 +6334,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", @@ -6327,7 +6374,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 22178f103..4293287d2 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, ) @@ -65,7 +66,9 @@ from langgraph.pregel import ( ) 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.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, @@ -209,6 +212,91 @@ async def test_node_cancellation_on_other_node_exception() -> None: assert inner_task_cancelled +async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_node_count = 0 + + async def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + raise NodeInterrupt("Just because...") + return {"my_key": " all good"} + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert await tool_two.ainvoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + tool_two = tool_two_graph.compile(checkpointer=saver) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke( + {"my_key": "value ⛰️", "market": "DE"}, thread1 + ) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("during", "Just because..."),), + ), + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + + @pytest.mark.parametrize( "checkpointer_name", ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], @@ -803,7 +891,7 @@ async def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 6, "writes": 5}, + metadata={"source": "loop", "step": 6, "writes": {"two": 5}}, created_at=AnyStr(), parent_config=history[1].config, ), @@ -818,7 +906,7 @@ async def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 5, "writes": None}, + metadata={"source": "loop", "step": 5, "writes": {"one": None}}, created_at=AnyStr(), parent_config=history[2].config, ), @@ -848,7 +936,7 @@ async def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 3, "writes": None}, + metadata={"source": "loop", "step": 3, "writes": {"one": None}}, created_at=AnyStr(), parent_config=history[4].config, ), @@ -878,7 +966,7 @@ async def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 1, "writes": 4}, + metadata={"source": "loop", "step": 1, "writes": {"two": 4}}, created_at=AnyStr(), parent_config=history[6].config, ), @@ -893,7 +981,7 @@ async def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "step": 0, "writes": None}, + metadata={"source": "loop", "step": 0, "writes": {"one": None}}, created_at=AnyStr(), parent_config=history[7].config, ), @@ -2217,12 +2305,14 @@ async def test_conditional_graph() -> None: "step": 0, "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } } }, }, @@ -2443,12 +2533,14 @@ async def test_conditional_graph() -> None: "step": 0, "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } } }, }, @@ -2675,12 +2767,14 @@ async def test_conditional_graph() -> None: "step": 0, "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } } }, }, @@ -4783,10 +4877,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 ) @@ -4803,14 +4920,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", @@ -4870,7 +4989,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", @@ -4918,7 +5037,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"],), + ] diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 13d5560db..ba8aa2f42 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -40,8 +40,14 @@ from langgraph_sdk.schema import ( logger = logging.getLogger(__name__) +RESERVED_HEADERS = ("x-api-key",) + + def get_client( - *, url: Optional[str] = None, api_key: Optional[str] = None + *, + url: Optional[str] = None, + api_key: Optional[str] = None, + headers: Optional[dict[str, str]] = None, ) -> LangGraphClient: """Get a LangGraphClient instance. @@ -53,6 +59,7 @@ def get_client( 2. LANGGRAPH_API_KEY 3. LANGSMITH_API_KEY 4. LANGCHAIN_API_KEY + headers: Optional custom headers """ transport: Optional[httpx.AsyncBaseTransport] = None if url is None: @@ -65,17 +72,12 @@ def get_client( url = "http://localhost:8123" if transport is None: transport = httpx.AsyncHTTPTransport(retries=5) - headers = { - "User-Agent": f"langgraph-sdk-py/{langgraph_sdk.__version__}", - } - api_key = _get_api_key(api_key) - if api_key: - headers["x-api-key"] = api_key + client = httpx.AsyncClient( base_url=url, transport=transport, timeout=httpx.Timeout(connect=5, read=60, write=60, pool=5), - headers=headers, + headers=_get_headers(api_key, headers), ) return LangGraphClient(client) @@ -1695,3 +1697,23 @@ def _get_api_key(api_key: Optional[str] = None) -> Optional[str]: if env := os.getenv(f"{prefix}_API_KEY"): return env.strip().strip('"').strip("'") return None # type: ignore + + +def _get_headers( + api_key: Optional[str], custom_headers: Optional[dict[str, str]] +) -> dict[str, str]: + """Combine api_key and custom user-provided headers.""" + custom_headers = custom_headers or {} + for header in RESERVED_HEADERS: + if header in custom_headers: + raise ValueError(f"Cannot set reserved header '{header}'") + + headers = { + "User-Agent": f"langgraph-sdk-py/{langgraph_sdk.__version__}", + **custom_headers, + } + api_key = _get_api_key(api_key) + if api_key: + headers["x-api-key"] = api_key + + return headers diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml index 3c50594db..2755c3fc6 100644 --- a/libs/sdk-py/pyproject.toml +++ b/libs/sdk-py/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-sdk" -version = "0.1.27" +version = "0.1.28" description = "SDK for interacting with LangGraph API" authors = [] license = "MIT"