from __future__ import annotations import asyncio import concurrent import concurrent.futures import contextlib import queue import warnings import weakref from collections import defaultdict, deque from collections.abc import ( AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence, ) from dataclasses import is_dataclass from functools import partial from inspect import isclass from typing import ( Any, Generic, cast, get_type_hints, ) from uuid import UUID, uuid5 from langchain_core.globals import get_debug from langchain_core.runnables import ( RunnableSequence, ) from langchain_core.runnables.base import Input, Output from langchain_core.runnables.config import ( RunnableConfig, get_async_callback_manager_for_config, get_callback_manager_for_config, ) from langchain_core.runnables.graph import Graph from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, CheckpointTuple, ) from langgraph.store.base import BaseStore from pydantic import BaseModel, TypeAdapter from typing_extensions import Self, Unpack, deprecated, is_typeddict from langgraph._internal._config import ( ensure_config, merge_configs, patch_checkpoint_map, patch_config, patch_configurable, recast_checkpoint_ns, ) from langgraph._internal._constants import ( CACHE_NS_WRITES, CONF, CONFIG_KEY_CACHE, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_DURABILITY, CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, CONFIG_KEY_RUNNER_SUBMIT, CONFIG_KEY_RUNTIME, CONFIG_KEY_SEND, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, CONFIG_KEY_THREAD_ID, ERROR, INPUT, INTERRUPT, NS_END, NS_SEP, NULL_TASK_ID, PUSH, TASKS, ) from langgraph._internal._pydantic import create_model from langgraph._internal._queue import ( # type: ignore[attr-defined] AsyncQueue, SyncQueue, ) from langgraph._internal._runnable import ( Runnable, RunnableLike, RunnableSeq, coerce_to_runnable, ) from langgraph._internal._typing import MISSING, DeprecatedKwargs from langgraph.channels.base import BaseChannel from langgraph.channels.topic import Topic from langgraph.config import get_config from langgraph.constants import END from langgraph.errors import ( ErrorCode, GraphRecursionError, InvalidUpdateError, create_error_message, ) from langgraph.managed.base import ManagedValueSpec from langgraph.pregel._algo import ( PregelTaskWrites, _scratchpad, apply_writes, local_read, prepare_next_tasks, ) from langgraph.pregel._call import identifier from langgraph.pregel._checkpoint import ( channels_from_checkpoint, copy_checkpoint, create_checkpoint, empty_checkpoint, ) from langgraph.pregel._draw import draw_graph from langgraph.pregel._io import map_input, read_channels from langgraph.pregel._loop import AsyncPregelLoop, SyncPregelLoop from langgraph.pregel._messages import StreamMessagesHandler from langgraph.pregel._read import DEFAULT_BOUND, PregelNode from langgraph.pregel._retry import RetryPolicy from langgraph.pregel._runner import PregelRunner 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.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtocol from langgraph.runtime import DEFAULT_RUNTIME, Runtime from langgraph.types import ( All, CachePolicy, Checkpointer, Command, Durability, Interrupt, Send, StateSnapshot, StateUpdate, StreamMode, ) from langgraph.typing import ContextT, InputT, OutputT, StateT from langgraph.warnings import LangGraphDeprecatedSinceV10 try: from langchain_core.tracers._streaming import _StreamingCallbackHandler except ImportError: _StreamingCallbackHandler = None # type: ignore __all__ = ("NodeBuilder", "Pregel") _WriteValue = Callable[[Input], Output] | Any class NodeBuilder: __slots__ = ( "_channels", "_triggers", "_tags", "_metadata", "_writes", "_bound", "_retry_policy", "_cache_policy", ) _channels: str | list[str] _triggers: list[str] _tags: list[str] _metadata: dict[str, Any] _writes: list[ChannelWriteEntry] _bound: Runnable _retry_policy: list[RetryPolicy] _cache_policy: CachePolicy | None def __init__( self, ) -> None: self._channels = [] self._triggers = [] self._tags = [] self._metadata = {} self._writes = [] self._bound = DEFAULT_BOUND self._retry_policy = [] self._cache_policy = None def subscribe_only( self, channel: str, ) -> Self: """Subscribe to a single channel.""" if not self._channels: self._channels = channel else: raise ValueError( "Cannot subscribe to single channels when other channels are already subscribed to" ) self._triggers.append(channel) return self def subscribe_to( self, *channels: str, read: bool = True, ) -> Self: """Add channels to subscribe to. Node will be invoked when any of these channels are updated, with a dict of the channel values as input. Args: channels: Channel name(s) to subscribe to read: If `True`, the channels will be included in the input to the node. Otherwise, they will trigger the node without being sent in input. Returns: Self for chaining """ if isinstance(self._channels, str): raise ValueError( "Cannot subscribe to channels when subscribed to a single channel" ) if read: if not self._channels: self._channels = list(channels) else: self._channels.extend(channels) if isinstance(channels, str): self._triggers.append(channels) else: self._triggers.extend(channels) return self def read_from( self, *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" assert isinstance(self._channels, list), ( "Cannot read additional channels when subscribed to single channels" ) self._channels.extend(channels) return self def do( self, node: RunnableLike, ) -> Self: """Adds the specified node.""" if self._bound is not DEFAULT_BOUND: self._bound = RunnableSeq( self._bound, coerce_to_runnable(node, name=None, trace=True) ) else: self._bound = coerce_to_runnable(node, name=None, trace=True) return self def write_to( self, *channels: str | ChannelWriteEntry, **kwargs: _WriteValue, ) -> Self: """Add channel writes. Args: *channels: Channel names to write to. **kwargs: Channel name and value mappings. Returns: Self for chaining """ self._writes.extend( ChannelWriteEntry(c) if isinstance(c, str) else c for c in channels ) self._writes.extend( ChannelWriteEntry(k, mapper=v) if callable(v) else ChannelWriteEntry(k, value=v) for k, v in kwargs.items() ) return self def meta(self, *tags: str, **metadata: Any) -> Self: """Add tags or metadata to the node.""" self._tags.extend(tags) self._metadata.update(metadata) return self def add_retry_policies(self, *policies: RetryPolicy) -> Self: """Adds retry policies to the node.""" self._retry_policy.extend(policies) return self def add_cache_policy(self, policy: CachePolicy) -> Self: """Adds cache policies to the node.""" self._cache_policy = policy return self def build(self) -> PregelNode: """Builds the node.""" return PregelNode( channels=self._channels, triggers=self._triggers, tags=self._tags, metadata=self._metadata, writers=[ChannelWrite(self._writes)], bound=self._bound, retry_policy=self._retry_policy, cache_policy=self._cache_policy, ) class Pregel( PregelProtocol[StateT, ContextT, InputT, OutputT], Generic[StateT, ContextT, InputT, OutputT], ): """Pregel manages the runtime behavior for LangGraph applications. ## Overview Pregel combines [**actors**](https://en.wikipedia.org/wiki/Actor_model) and **channels** into a single application. **Actors** read data from channels and write data to channels. Pregel organizes the execution of the application into multiple steps, following the **Pregel Algorithm**/**Bulk Synchronous Parallel** model. Each step consists of three phases: - **Plan**: Determine which **actors** to execute in this step. For example, in the first step, select the **actors** that subscribe to the special **input** channels; in subsequent steps, select the **actors** that subscribe to channels updated in the previous step. - **Execution**: Execute all selected **actors** in parallel, until all complete, or one fails, or a timeout is reached. During this phase, channel updates are invisible to actors until the next step. - **Update**: Update the channels with the values written by the **actors** in this step. Repeat until no **actors** are selected for execution, or a maximum number of steps is reached. ## Actors An **actor** is a `PregelNode`. It subscribes to channels, reads data from them, and writes data to them. It can be thought of as an **actor** in the Pregel algorithm. `PregelNodes` implement LangChain's Runnable interface. ## Channels Channels are used to communicate between actors (`PregelNodes`). Each channel has a value type, an update type, and an update function – which takes a sequence of updates and modifies the stored value. Channels can be used to send data from one chain to another, or to send data from a chain to itself in a future step. LangGraph provides a number of built-in channels: ### Basic channels: LastValue and Topic - `LastValue`: The default channel, stores the last value sent to the channel, useful for input and output values, or for sending data from one step to the next - `Topic`: A configurable PubSub Topic, useful for sending multiple values between *actors*, or for accumulating output. Can be configured to deduplicate values, and/or to accumulate values over the course of multiple steps. ### Advanced channels: Context and BinaryOperatorAggregate - `Context`: exposes the value of a context manager, managing its lifecycle. Useful for accessing external resources that require setup and/or teardown. e.g. `client = Context(httpx.Client)` - `BinaryOperatorAggregate`: stores a persistent value, updated by applying a binary operator to the current value and each update sent to the channel, useful for computing aggregates over multiple steps. e.g. `total = BinaryOperatorAggregate(int, operator.add)` ## Examples Most users will interact with Pregel via a [StateGraph (Graph API)][langgraph.graph.StateGraph] or via an [entrypoint (Functional API)][langgraph.func.entrypoint]. However, for **advanced** use cases, Pregel can be used directly. If you're not sure whether you need to use Pregel directly, then the answer is probably no - you should use the Graph API or Functional API instead. These are higher-level interfaces that will compile down to Pregel under the hood. Here are some examples to give you a sense of how it works: Example: Single node application ```python from langgraph.channels import EphemeralValue from langgraph.pregel import Pregel, NodeBuilder node1 = ( NodeBuilder().subscribe_only("a") .do(lambda x: x + x) .write_to("b") ) app = Pregel( nodes={"node1": node1}, channels={ "a": EphemeralValue(str), "b": EphemeralValue(str), }, input_channels=["a"], output_channels=["b"], ) app.invoke({"a": "foo"}) ``` ```con {'b': 'foofoo'} ``` Example: Using multiple nodes and multiple output channels ```python from langgraph.channels import LastValue, EphemeralValue from langgraph.pregel import Pregel, NodeBuilder node1 = ( NodeBuilder().subscribe_only("a") .do(lambda x: x + x) .write_to("b") ) node2 = ( NodeBuilder().subscribe_to("b") .do(lambda x: x["b"] + x["b"]) .write_to("c") ) app = Pregel( nodes={"node1": node1, "node2": node2}, channels={ "a": EphemeralValue(str), "b": LastValue(str), "c": EphemeralValue(str), }, input_channels=["a"], output_channels=["b", "c"], ) app.invoke({"a": "foo"}) ``` ```con {'b': 'foofoo', 'c': 'foofoofoofoo'} ``` Example: Using a Topic channel ```python from langgraph.channels import LastValue, EphemeralValue, Topic from langgraph.pregel import Pregel, NodeBuilder node1 = ( NodeBuilder().subscribe_only("a") .do(lambda x: x + x) .write_to("b", "c") ) node2 = ( NodeBuilder().subscribe_only("b") .do(lambda x: x + x) .write_to("c") ) app = Pregel( nodes={"node1": node1, "node2": node2}, channels={ "a": EphemeralValue(str), "b": EphemeralValue(str), "c": Topic(str, accumulate=True), }, input_channels=["a"], output_channels=["c"], ) app.invoke({"a": "foo"}) ``` ```pycon {"c": ["foofoo", "foofoofoofoo"]} ``` Example: Using a `BinaryOperatorAggregate` channel ```python from langgraph.channels import EphemeralValue, BinaryOperatorAggregate from langgraph.pregel import Pregel, NodeBuilder node1 = ( NodeBuilder().subscribe_only("a") .do(lambda x: x + x) .write_to("b", "c") ) node2 = ( NodeBuilder().subscribe_only("b") .do(lambda x: x + x) .write_to("c") ) def reducer(current, update): if current: return current + " | " + update else: return update app = Pregel( nodes={"node1": node1, "node2": node2}, channels={ "a": EphemeralValue(str), "b": EphemeralValue(str), "c": BinaryOperatorAggregate(str, operator=reducer), }, input_channels=["a"], output_channels=["c"], ) app.invoke({"a": "foo"}) ``` ```con {'c': 'foofoo | foofoofoofoo'} ``` Example: Introducing a cycle This example demonstrates how to introduce a cycle in the graph, by having a chain write to a channel it subscribes to. Execution will continue until a `None` value is written to the channel. ```python from langgraph.channels import EphemeralValue from langgraph.pregel import Pregel, NodeBuilder, ChannelWriteEntry example_node = ( NodeBuilder() .subscribe_only("value") .do(lambda x: x + x if len(x) < 10 else None) .write_to(ChannelWriteEntry(channel="value", skip_none=True)) ) app = Pregel( nodes={"example_node": example_node}, channels={ "value": EphemeralValue(str), }, input_channels=["value"], output_channels=["value"], ) app.invoke({"value": "a"}) ``` ```con {'value': 'aaaaaaaaaaaaaaaa'} ``` """ nodes: dict[str, PregelNode] channels: dict[str, BaseChannel | ManagedValueSpec] stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" stream_eager: bool = False """Whether to force emitting stream events eagerly, automatically turned on for stream_mode "messages" and "custom".""" output_channels: str | Sequence[str] stream_channels: str | Sequence[str] | None = None """Channels to stream, defaults to all channels not in reserved channels""" interrupt_after_nodes: All | Sequence[str] interrupt_before_nodes: All | Sequence[str] input_channels: str | Sequence[str] step_timeout: float | None = None """Maximum time to wait for a step to complete, in seconds.""" debug: bool """Whether to print debug information during execution.""" checkpointer: Checkpointer = None """`Checkpointer` used to save and load graph state.""" store: BaseStore | None = None """Memory store to use for SharedValues.""" cache: BaseCache | None = None """Cache to use for storing node results.""" retry_policy: Sequence[RetryPolicy] = () """Retry policies to use when running tasks. Empty set disables retries.""" cache_policy: CachePolicy | None = None """Cache policy to use for all nodes. Can be overridden by individual nodes.""" context_schema: type[ContextT] | None = None """Specifies the schema for the context object that will be passed to the workflow.""" config: RunnableConfig | None = None name: str = "LangGraph" trigger_to_nodes: Mapping[str, Sequence[str]] def __init__( self, *, nodes: dict[str, PregelNode | NodeBuilder], channels: dict[str, BaseChannel | ManagedValueSpec] | None, auto_validate: bool = True, stream_mode: StreamMode = "values", stream_eager: bool = False, output_channels: str | Sequence[str], stream_channels: str | Sequence[str] | None = None, interrupt_after_nodes: All | Sequence[str] = (), interrupt_before_nodes: All | Sequence[str] = (), input_channels: str | Sequence[str], step_timeout: float | None = None, debug: bool | None = None, checkpointer: BaseCheckpointSaver | None = None, store: BaseStore | None = None, cache: BaseCache | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, context_schema: type[ContextT] | None = None, config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", **deprecated_kwargs: Unpack[DeprecatedKwargs], ) -> None: if ( config_type := deprecated_kwargs.get("config_type", MISSING) ) is not MISSING: warnings.warn( "`config_type` is deprecated and will be removed. Please use `context_schema` instead.", category=LangGraphDeprecatedSinceV10, stacklevel=2, ) if context_schema is None: context_schema = cast(type[ContextT], config_type) self.nodes = { k: v.build() if isinstance(v, NodeBuilder) else v for k, v in nodes.items() } self.channels = channels or {} if TASKS in self.channels and not isinstance(self.channels[TASKS], Topic): raise ValueError( f"Channel '{TASKS}' is reserved and cannot be used in the graph." ) else: self.channels[TASKS] = Topic(Send, accumulate=False) self.stream_mode = stream_mode self.stream_eager = stream_eager self.output_channels = output_channels self.stream_channels = stream_channels self.interrupt_after_nodes = interrupt_after_nodes self.interrupt_before_nodes = interrupt_before_nodes self.input_channels = input_channels self.step_timeout = step_timeout self.debug = debug if debug is not None else get_debug() self.checkpointer = checkpointer self.store = store self.cache = cache self.retry_policy = ( (retry_policy,) if isinstance(retry_policy, RetryPolicy) else retry_policy ) self.cache_policy = cache_policy self.context_schema = context_schema self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name if auto_validate: self.validate() def get_graph( self, config: RunnableConfig | None = None, *, xray: int | bool = False ) -> Graph: """Return a drawable representation of the computation graph.""" # gather subgraphs if xray: subgraphs = { k: v.get_graph( config, xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1, ) for k, v in self.get_subgraphs() } else: subgraphs = {} return draw_graph( merge_configs(self.config, config), nodes=self.nodes, specs=self.channels, input_channels=self.input_channels, interrupt_after_nodes=self.interrupt_after_nodes, interrupt_before_nodes=self.interrupt_before_nodes, trigger_to_nodes=self.trigger_to_nodes, checkpointer=self.checkpointer, subgraphs=subgraphs, ) async def aget_graph( self, config: RunnableConfig | None = None, *, xray: int | bool = False ) -> Graph: """Return a drawable representation of the computation graph.""" # gather subgraphs if xray: subpregels: dict[str, PregelProtocol] = { k: v async for k, v in self.aget_subgraphs() } subgraphs = { k: v for k, v in zip( subpregels, await asyncio.gather( *( p.aget_graph( config, xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1, ) for p in subpregels.values() ) ), ) } else: subgraphs = {} return draw_graph( merge_configs(self.config, config), nodes=self.nodes, specs=self.channels, input_channels=self.input_channels, interrupt_after_nodes=self.interrupt_after_nodes, interrupt_before_nodes=self.interrupt_before_nodes, trigger_to_nodes=self.trigger_to_nodes, checkpointer=self.checkpointer, subgraphs=subgraphs, ) def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]: """Mime bundle used by Jupyter to display the graph""" return { "text/plain": repr(self), "image/png": self.get_graph().draw_mermaid_png(), } def copy(self, update: dict[str, Any] | None = None) -> Self: attrs = {k: v for k, v in self.__dict__.items() if k != "__orig_class__"} attrs.update(update or {}) return self.__class__(**attrs) def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: """Create a copy of the Pregel object with an updated config.""" return self.copy( {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))} ) def validate(self) -> Self: validate_graph( self.nodes, {k: v for k, v in self.channels.items() if isinstance(v, BaseChannel)}, {k: v for k, v in self.channels.items() if not isinstance(v, BaseChannel)}, self.input_channels, self.output_channels, self.stream_channels, self.interrupt_after_nodes, self.interrupt_before_nodes, ) self.trigger_to_nodes = _trigger_to_nodes(self.nodes) return self @deprecated( "`config_schema` is deprecated. Use `get_context_jsonschema` for the relevant schema instead.", category=None, ) def config_schema(self, *, include: Sequence[str] | None = None) -> type[BaseModel]: warnings.warn( "`config_schema` is deprecated. Use `get_context_jsonschema` for the relevant schema instead.", category=LangGraphDeprecatedSinceV10, stacklevel=2, ) include = include or [] fields = { **( {"configurable": (self.context_schema, None)} if self.context_schema else {} ), **{ field_name: (field_type, None) for field_name, field_type in get_type_hints(RunnableConfig).items() if field_name in [i for i in include if i != "configurable"] }, } return create_model(self.get_name("Config"), field_definitions=fields) @deprecated( "`get_config_jsonschema` is deprecated. Use `get_context_jsonschema` instead.", category=None, ) def get_config_jsonschema( self, *, include: Sequence[str] | None = None ) -> dict[str, Any]: warnings.warn( "`get_config_jsonschema` is deprecated. Use `get_context_jsonschema` instead.", category=LangGraphDeprecatedSinceV10, stacklevel=2, ) with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=LangGraphDeprecatedSinceV10) schema = self.config_schema(include=include) return schema.model_json_schema() def get_context_jsonschema(self) -> dict[str, Any] | None: if (context_schema := self.context_schema) is None: return None if isclass(context_schema) and issubclass(context_schema, BaseModel): return context_schema.model_json_schema() elif is_typeddict(context_schema) or is_dataclass(context_schema): return TypeAdapter(context_schema).json_schema() else: raise ValueError( f"Invalid context schema type: {context_schema}. Must be a BaseModel, TypedDict or dataclass." ) @property def InputType(self) -> Any: if isinstance(self.input_channels, str): channel = self.channels[self.input_channels] if isinstance(channel, BaseChannel): return channel.UpdateType def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]: config = merge_configs(self.config, config) if isinstance(self.input_channels, str): return super().get_input_schema(config) else: return create_model( self.get_name("Input"), field_definitions={ k: (c.UpdateType, None) for k in self.input_channels or self.channels.keys() if (c := self.channels[k]) and isinstance(c, BaseChannel) }, ) def get_input_jsonschema( self, config: RunnableConfig | None = None ) -> dict[str, Any]: schema = self.get_input_schema(config) return schema.model_json_schema() @property def OutputType(self) -> Any: if isinstance(self.output_channels, str): channel = self.channels[self.output_channels] if isinstance(channel, BaseChannel): return channel.ValueType def get_output_schema( self, config: RunnableConfig | None = None ) -> type[BaseModel]: config = merge_configs(self.config, config) if isinstance(self.output_channels, str): return super().get_output_schema(config) else: return create_model( self.get_name("Output"), field_definitions={ k: (c.ValueType, None) for k in self.output_channels if (c := self.channels[k]) and isinstance(c, BaseChannel) }, ) def get_output_jsonschema( self, config: RunnableConfig | None = None ) -> dict[str, Any]: schema = self.get_output_schema(config) return schema.model_json_schema() @property def stream_channels_list(self) -> Sequence[str]: stream_channels = self.stream_channels_asis return ( [stream_channels] if isinstance(stream_channels, str) else stream_channels ) @property def stream_channels_asis(self) -> str | Sequence[str]: return self.stream_channels or [ k for k in self.channels if isinstance(self.channels[k], BaseChannel) ] def get_subgraphs( self, *, namespace: str | None = None, recurse: bool = False ) -> Iterator[tuple[str, PregelProtocol]]: """Get the subgraphs of the graph. Args: namespace: The namespace to filter the subgraphs by. recurse: Whether to recurse into the subgraphs. If `False`, only the immediate subgraphs will be returned. Returns: An iterator of the `(namespace, subgraph)` pairs. """ for name, node in self.nodes.items(): # filter by prefix if namespace is not None: if not namespace.startswith(name): continue # find the subgraph, if any graph = node.subgraphs[0] if node.subgraphs else None # if found, yield recursively if graph: if name == namespace: yield name, graph return # we found it, stop searching if namespace is None: yield name, graph if recurse and isinstance(graph, Pregel): if namespace is not None: namespace = namespace[len(name) + 1 :] yield from ( (f"{name}{NS_SEP}{n}", s) for n, s in graph.get_subgraphs( namespace=namespace, recurse=recurse ) ) async def aget_subgraphs( self, *, namespace: str | None = None, recurse: bool = False ) -> AsyncIterator[tuple[str, PregelProtocol]]: """Get the subgraphs of the graph. Args: namespace: The namespace to filter the subgraphs by. recurse: Whether to recurse into the subgraphs. If `False`, only the immediate subgraphs will be returned. Returns: An iterator of the `(namespace, subgraph)` pairs. """ for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse): yield name, node def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None: """Migrate a saved checkpoint to new channel layout.""" if checkpoint["v"] < 4 and checkpoint.get("pending_sends"): pending_sends: list[Send] = checkpoint.pop("pending_sends") checkpoint["channel_values"][TASKS] = pending_sends checkpoint["channel_versions"][TASKS] = max( checkpoint["channel_versions"].values() ) def _prepare_state_snapshot( self, config: RunnableConfig, saved: CheckpointTuple | None, recurse: BaseCheckpointSaver | None = None, apply_pending_writes: bool = False, ) -> StateSnapshot: if not saved: return StateSnapshot( values={}, next=(), config=config, metadata=None, created_at=None, parent_config=None, tasks=(), interrupts=(), ) # migrate checkpoint if needed self._migrate_checkpoint(saved.checkpoint) step = saved.metadata.get("step", -1) + 1 stop = step + 2 channels, managed = channels_from_checkpoint( self.channels, saved.checkpoint, ) # tasks for this checkpoint next_tasks = prepare_next_tasks( saved.checkpoint, saved.pending_writes or [], self.nodes, channels, managed, saved.config, step, stop, for_execution=True, store=self.store, checkpointer=( self.checkpointer if isinstance(self.checkpointer, BaseCheckpointSaver) else None ), manager=None, ) # get the subgraphs subgraphs = dict(self.get_subgraphs()) parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") task_states: dict[str, RunnableConfig | StateSnapshot] = {} for task in next_tasks.values(): if task.name not in subgraphs: continue # assemble checkpoint_ns for this task task_ns = f"{task.name}{NS_END}{task.id}" if parent_ns: task_ns = f"{parent_ns}{NS_SEP}{task_ns}" if not recurse: # set config as signal that subgraph checkpoints exist config = { CONF: { "thread_id": saved.config[CONF]["thread_id"], CONFIG_KEY_CHECKPOINT_NS: task_ns, } } task_states[task.id] = config else: # get the state of the subgraph config = { CONF: { CONFIG_KEY_CHECKPOINTER: recurse, "thread_id": saved.config[CONF]["thread_id"], CONFIG_KEY_CHECKPOINT_NS: task_ns, } } task_states[task.id] = subgraphs[task.name].get_state( config, subgraphs=True ) # apply pending writes if null_writes := [ w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: apply_writes( saved.checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, self.trigger_to_nodes, ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes( saved.checkpoint, channels, tasks, None, self.trigger_to_nodes ) tasks_with_writes = tasks_w_writes( next_tasks.values(), saved.pending_writes, task_states, self.stream_channels_asis, ) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), tuple(t.name for t in next_tasks.values() if not t.writes), patch_checkpoint_map(saved.config, saved.metadata), saved.metadata, saved.checkpoint["ts"], patch_checkpoint_map(saved.parent_config, saved.metadata), tasks_with_writes, tuple([i for task in tasks_with_writes for i in task.interrupts]), ) async def _aprepare_state_snapshot( self, config: RunnableConfig, saved: CheckpointTuple | None, recurse: BaseCheckpointSaver | None = None, apply_pending_writes: bool = False, ) -> StateSnapshot: if not saved: return StateSnapshot( values={}, next=(), config=config, metadata=None, created_at=None, parent_config=None, tasks=(), interrupts=(), ) # migrate checkpoint if needed self._migrate_checkpoint(saved.checkpoint) step = saved.metadata.get("step", -1) + 1 stop = step + 2 channels, managed = channels_from_checkpoint( self.channels, saved.checkpoint, ) # tasks for this checkpoint next_tasks = prepare_next_tasks( saved.checkpoint, saved.pending_writes or [], self.nodes, channels, managed, saved.config, step, stop, for_execution=True, store=self.store, checkpointer=( self.checkpointer if isinstance(self.checkpointer, BaseCheckpointSaver) else None ), manager=None, ) # get the subgraphs subgraphs = {n: g async for n, g in self.aget_subgraphs()} parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") task_states: dict[str, RunnableConfig | StateSnapshot] = {} for task in next_tasks.values(): if task.name not in subgraphs: continue # assemble checkpoint_ns for this task task_ns = f"{task.name}{NS_END}{task.id}" if parent_ns: task_ns = f"{parent_ns}{NS_SEP}{task_ns}" if not recurse: # set config as signal that subgraph checkpoints exist config = { CONF: { "thread_id": saved.config[CONF]["thread_id"], CONFIG_KEY_CHECKPOINT_NS: task_ns, } } task_states[task.id] = config else: # get the state of the subgraph config = { CONF: { CONFIG_KEY_CHECKPOINTER: recurse, "thread_id": saved.config[CONF]["thread_id"], CONFIG_KEY_CHECKPOINT_NS: task_ns, } } task_states[task.id] = await subgraphs[task.name].aget_state( config, subgraphs=True ) # apply pending writes if null_writes := [ w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: apply_writes( saved.checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], None, self.trigger_to_nodes, ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes( saved.checkpoint, channels, tasks, None, self.trigger_to_nodes ) tasks_with_writes = tasks_w_writes( next_tasks.values(), saved.pending_writes, task_states, self.stream_channels_asis, ) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), tuple(t.name for t in next_tasks.values() if not t.writes), patch_checkpoint_map(saved.config, saved.metadata), saved.metadata, saved.checkpoint["ts"], patch_checkpoint_map(saved.parent_config, saved.metadata), tasks_with_writes, tuple([i for task in tasks_with_writes for i in task.interrupts]), ) def get_state( self, config: RunnableConfig, *, subgraphs: bool = False ) -> StateSnapshot: """Get the current state of the graph.""" checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): return pregel.get_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), subgraphs=subgraphs, ) else: raise ValueError(f"Subgraph {recast} not found") config = merge_configs(self.config, config) if self.config else config if self.checkpointer is True: ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) config = merge_configs( config, {CONF: {CONFIG_KEY_CHECKPOINT_NS: recast_checkpoint_ns(ns)}} ) thread_id = config[CONF][CONFIG_KEY_THREAD_ID] if not isinstance(thread_id, str): config[CONF][CONFIG_KEY_THREAD_ID] = str(thread_id) saved = checkpointer.get_tuple(config) return self._prepare_state_snapshot( config, saved, recurse=checkpointer if subgraphs else None, apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF], ) async def aget_state( self, config: RunnableConfig, *, subgraphs: bool = False ) -> StateSnapshot: """Get the current state of the graph.""" checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): return await pregel.aget_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), subgraphs=subgraphs, ) else: raise ValueError(f"Subgraph {recast} not found") config = merge_configs(self.config, config) if self.config else config if self.checkpointer is True: ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) config = merge_configs( config, {CONF: {CONFIG_KEY_CHECKPOINT_NS: recast_checkpoint_ns(ns)}} ) thread_id = config[CONF][CONFIG_KEY_THREAD_ID] if not isinstance(thread_id, str): config[CONF][CONFIG_KEY_THREAD_ID] = str(thread_id) saved = await checkpointer.aget_tuple(config) return await self._aprepare_state_snapshot( config, saved, recurse=checkpointer if subgraphs else None, apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF], ) def get_state_history( self, config: RunnableConfig, *, filter: dict[str, Any] | None = None, before: RunnableConfig | None = None, limit: int | None = None, ) -> Iterator[StateSnapshot]: """Get the history of the state of the graph.""" config = ensure_config(config) checkpointer: BaseCheckpointSaver | None = config[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): yield from pregel.get_state_history( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), filter=filter, before=before, limit=limit, ) return else: raise ValueError(f"Subgraph {recast} not found") config = merge_configs( self.config, config, { CONF: { CONFIG_KEY_CHECKPOINT_NS: checkpoint_ns, CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID]), } }, ) # eagerly consume list() to avoid holding up the db cursor for checkpoint_tuple in list( checkpointer.list(config, before=before, limit=limit, filter=filter) ): yield self._prepare_state_snapshot( checkpoint_tuple.config, checkpoint_tuple ) async def aget_state_history( self, config: RunnableConfig, *, filter: dict[str, Any] | None = None, before: RunnableConfig | None = None, limit: int | None = None, ) -> AsyncIterator[StateSnapshot]: """Asynchronously get the history of the state of the graph.""" config = ensure_config(config) checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): async for state in pregel.aget_state_history( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), filter=filter, before=before, limit=limit, ): yield state return else: raise ValueError(f"Subgraph {recast} not found") config = merge_configs( self.config, config, { CONF: { CONFIG_KEY_CHECKPOINT_NS: checkpoint_ns, CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID]), } }, ) # eagerly consume list() to avoid holding up the db cursor for checkpoint_tuple in [ c async for c in checkpointer.alist( config, before=before, limit=limit, filter=filter ) ]: yield await self._aprepare_state_snapshot( checkpoint_tuple.config, checkpoint_tuple ) def bulk_update_state( self, config: RunnableConfig, supersteps: Sequence[Sequence[StateUpdate]], ) -> RunnableConfig: """Apply updates to the graph state in bulk. Requires a checkpointer to be set. Args: config: The config to apply the updates to. supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state. Each update is a tuple of the form `(values, as_node, task_id)` where `task_id` is optional. Raises: ValueError: If no checkpointer is set or no updates are provided. InvalidUpdateError: If an invalid update is provided. Returns: RunnableConfig: The updated config. """ checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") if len(supersteps) == 0: raise ValueError("No supersteps provided") if any(len(u) == 0 for u in supersteps): raise ValueError("No updates provided") # delegate to subgraph if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): return pregel.bulk_update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), supersteps, ) else: raise ValueError(f"Subgraph {recast} not found") def perform_superstep( input_config: RunnableConfig, updates: Sequence[StateUpdate] ) -> RunnableConfig: # get last checkpoint config = ensure_config(self.config, input_config) saved = checkpointer.get_tuple(config) if saved is not None: self._migrate_checkpoint(saved.checkpoint) checkpoint = ( copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() ) checkpoint_previous_versions = ( saved.checkpoint["channel_versions"].copy() if saved else {} ) step = saved.metadata.get("step", -1) if saved else -1 # merge configurable fields with previous checkpoint config checkpoint_config = patch_configurable( config, { CONFIG_KEY_CHECKPOINT_NS: config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" ) }, ) if saved: checkpoint_config = patch_configurable(config, saved.config[CONF]) channels, managed = channels_from_checkpoint( self.channels, checkpoint, ) values, as_node = updates[0][:2] # no values as END, just clear all tasks if values is None and as_node == END: if len(updates) > 1: raise InvalidUpdateError( "Cannot apply multiple updates when clearing state" ) if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, saved.pending_writes or [], self.nodes, channels, managed, saved.config, step + 1, step + 3, for_execution=True, store=self.store, checkpointer=checkpointer, manager=None, ) # apply null writes if null_writes := [ w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], checkpointer.get_next_version, self.trigger_to_nodes, ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) # clear all current tasks apply_writes( checkpoint, channels, next_tasks.values(), checkpointer.get_next_version, self.trigger_to_nodes, ) # save checkpoint next_config = checkpointer.put( checkpoint_config, create_checkpoint(checkpoint, channels, step), { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"], ), ) return patch_checkpoint_map( next_config, saved.metadata if saved else None ) # act as an input if as_node == INPUT: if len(updates) > 1: raise InvalidUpdateError( "Cannot apply multiple updates when updating as input" ) if input_writes := deque(map_input(self.input_channels, values)): apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, input_writes, [])], checkpointer.get_next_version, self.trigger_to_nodes, ) # apply input write to channels next_step = ( step + 1 if saved and saved.metadata.get("step") is not None else -1 ) next_config = checkpointer.put( checkpoint_config, create_checkpoint(checkpoint, channels, next_step), { "source": "input", "step": next_step, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"], ), ) # store the writes checkpointer.put_writes( next_config, input_writes, str(uuid5(UUID(checkpoint["id"]), INPUT)), ) return patch_checkpoint_map( next_config, saved.metadata if saved else None ) else: raise InvalidUpdateError( f"Received no input writes for {self.input_channels}" ) # copy checkpoint if as_node == "__copy__": if len(updates) > 1: raise InvalidUpdateError( "Cannot copy checkpoint with multiple updates" ) if saved is None: raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = checkpointer.put( saved.parent_config or patch_configurable( saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} ), next_checkpoint, { "source": "fork", "step": step + 1, "parents": saved.metadata.get("parents", {}), }, {}, ) # we want to both clone a checkpoint and update state in one go. # reuse the same task ID if possible. if isinstance(values, list) and len(values) > 0: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, saved.pending_writes or [], self.nodes, channels, managed, next_config, step + 2, step + 4, for_execution=True, store=self.store, checkpointer=checkpointer, manager=None, ) tasks_group_by = defaultdict(list) user_group_by: dict[str, list[StateUpdate]] = defaultdict(list) for task in next_tasks.values(): tasks_group_by[task.name].append(task.id) for item in values: if not isinstance(item, Sequence): raise InvalidUpdateError( f"Invalid update item: {item} when copying checkpoint" ) values, as_node = item[:2] user_group = user_group_by[as_node] tasks_group = tasks_group_by[as_node] target_idx = len(user_group) task_id = ( tasks_group[target_idx] if target_idx < len(tasks_group) else None ) user_group_by[as_node].append( StateUpdate(values=values, as_node=as_node, task_id=task_id) ) return perform_superstep( patch_checkpoint_map(next_config, saved.metadata), [item for lst in user_group_by.values() for item in lst], ) return patch_checkpoint_map(next_config, saved.metadata) # task ids can be provided in the StateUpdate, but if not, # we use the task id generated by prepare_next_tasks node_to_task_ids: dict[str, deque[str]] = defaultdict(deque) if saved is not None and saved.pending_writes is not None: # we call prepare_next_tasks to discover the task IDs that # would have been generated, so we can reuse them and # properly populate task.result in state history next_tasks = prepare_next_tasks( checkpoint, saved.pending_writes, self.nodes, channels, managed, saved.config, step + 1, step + 3, for_execution=True, store=self.store, checkpointer=checkpointer, manager=None, ) # collect task ids to reuse so we can properly attach task results for t in next_tasks.values(): node_to_task_ids[t.name].append(t.id) valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = [] if len(updates) == 1: values, as_node, task_id = updates[0] # find last node that updated the state, if not provided if as_node is None and len(self.nodes) == 1: as_node = tuple(self.nodes)[0] elif as_node is None and not any( v for vv in checkpoint["versions_seen"].values() for v in vv.values() ): if ( isinstance(self.input_channels, str) and self.input_channels in self.nodes ): as_node = self.input_channels elif as_node is None: last_seen_by_node = sorted( (v, n) for n, seen in checkpoint["versions_seen"].items() if n in self.nodes for v in seen.values() ) # if two nodes updated the state at the same time, it's ambiguous if last_seen_by_node: if len(last_seen_by_node) == 1: as_node = last_seen_by_node[0][1] elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: as_node = last_seen_by_node[-1][1] if as_node is None: raise InvalidUpdateError("Ambiguous update, specify as_node") if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") valid_updates.append((as_node, values, task_id)) else: for values, as_node, task_id in updates: if as_node is None: raise InvalidUpdateError( "as_node is required when applying multiple updates" ) if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") valid_updates.append((as_node, values, task_id)) run_tasks: list[PregelTaskWrites] = [] run_task_ids: list[str] = [] for as_node, values, provided_task_id in valid_updates: # create task to run all writers of the chosen node writers = self.nodes[as_node].flat_writers if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") writes: deque[tuple[str, Any]] = deque() task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) # get the task ids that were prepared for this node # if a task id was provided in the StateUpdate, we use it # otherwise, we use the next available task id prepared_task_ids = node_to_task_ids.get(as_node, deque()) task_id = provided_task_id or ( prepared_task_ids.popleft() if prepared_task_ids else str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) ) run_tasks.append(task) run_task_ids.append(task_id) run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] # execute task run.invoke( values, patch_config( config, run_name=self.name + "UpdateState", configurable={ # deque.extend is thread-safe CONFIG_KEY_SEND: writes.extend, CONFIG_KEY_TASK_ID: task_id, CONFIG_KEY_READ: partial( local_read, _scratchpad( None, [], task_id, "", None, step, step + 2, ), channels, managed, task, ), }, ), ) # save task writes for task_id, task in zip(run_task_ids, run_tasks): # channel writes are saved to current checkpoint channel_writes = [w for w in task.writes if w[0] != PUSH] if saved and channel_writes: checkpointer.put_writes(checkpoint_config, channel_writes, task_id) # apply to checkpoint and save apply_writes( checkpoint, channels, run_tasks, checkpointer.get_next_version, self.trigger_to_nodes, ) checkpoint = create_checkpoint(checkpoint, channels, step + 1) next_config = checkpointer.put( checkpoint_config, checkpoint, { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) for task_id, task in zip(run_task_ids, run_tasks): # save push writes if push_writes := [w for w in task.writes if w[0] == PUSH]: checkpointer.put_writes(next_config, push_writes, task_id) return patch_checkpoint_map(next_config, saved.metadata if saved else None) current_config = patch_configurable( config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])} ) for superstep in supersteps: current_config = perform_superstep(current_config, superstep) return current_config async def abulk_update_state( self, config: RunnableConfig, supersteps: Sequence[Sequence[StateUpdate]], ) -> RunnableConfig: """Asynchronously apply updates to the graph state in bulk. Requires a checkpointer to be set. Args: config: The config to apply the updates to. supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state. Each update is a tuple of the form `(values, as_node, task_id)` where `task_id` is optional. Raises: ValueError: If no checkpointer is set or no updates are provided. InvalidUpdateError: If an invalid update is provided. Returns: RunnableConfig: The updated config. """ checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") if len(supersteps) == 0: raise ValueError("No supersteps provided") if any(len(u) == 0 for u in supersteps): raise ValueError("No updates provided") # delegate to subgraph if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): return await pregel.abulk_update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), supersteps, ) else: raise ValueError(f"Subgraph {recast} not found") async def aperform_superstep( input_config: RunnableConfig, updates: Sequence[StateUpdate] ) -> RunnableConfig: # get last checkpoint config = ensure_config(self.config, input_config) saved = await checkpointer.aget_tuple(config) if saved is not None: self._migrate_checkpoint(saved.checkpoint) checkpoint = ( copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() ) checkpoint_previous_versions = ( saved.checkpoint["channel_versions"].copy() if saved else {} ) step = saved.metadata.get("step", -1) if saved else -1 # merge configurable fields with previous checkpoint config checkpoint_config = patch_configurable( config, { CONFIG_KEY_CHECKPOINT_NS: config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" ) }, ) if saved: checkpoint_config = patch_configurable(config, saved.config[CONF]) channels, managed = channels_from_checkpoint( self.channels, checkpoint, ) values, as_node = updates[0][:2] # no values, just clear all tasks if values is None and as_node == END: if len(updates) > 1: raise InvalidUpdateError( "Cannot apply multiple updates when clearing state" ) if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, saved.pending_writes or [], self.nodes, channels, managed, saved.config, step + 1, step + 3, for_execution=True, store=self.store, checkpointer=checkpointer, manager=None, ) # apply null writes if null_writes := [ w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID ]: apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, null_writes, [])], checkpointer.get_next_version, self.trigger_to_nodes, ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: if k in (ERROR, INTERRUPT): continue if tid not in next_tasks: continue next_tasks[tid].writes.append((k, v)) # clear all current tasks apply_writes( checkpoint, channels, next_tasks.values(), checkpointer.get_next_version, self.trigger_to_nodes, ) # save checkpoint next_config = await checkpointer.aput( checkpoint_config, create_checkpoint(checkpoint, channels, step), { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) return patch_checkpoint_map( next_config, saved.metadata if saved else None ) # act as an input if as_node == INPUT: if len(updates) > 1: raise InvalidUpdateError( "Cannot apply multiple updates when updating as input" ) if input_writes := deque(map_input(self.input_channels, values)): apply_writes( checkpoint, channels, [PregelTaskWrites((), INPUT, input_writes, [])], checkpointer.get_next_version, self.trigger_to_nodes, ) # apply input write to channels next_step = ( step + 1 if saved and saved.metadata.get("step") is not None else -1 ) next_config = await checkpointer.aput( checkpoint_config, create_checkpoint(checkpoint, channels, next_step), { "source": "input", "step": next_step, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"], ), ) # store the writes await checkpointer.aput_writes( next_config, input_writes, str(uuid5(UUID(checkpoint["id"]), INPUT)), ) return patch_checkpoint_map( next_config, saved.metadata if saved else None ) else: raise InvalidUpdateError( f"Received no input writes for {self.input_channels}" ) # no values, copy checkpoint if as_node == "__copy__": if len(updates) > 1: raise InvalidUpdateError( "Cannot copy checkpoint with multiple updates" ) if saved is None: raise InvalidUpdateError("Cannot copy a non-existent checkpoint") next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = await checkpointer.aput( saved.parent_config or patch_configurable( saved.config, {CONFIG_KEY_CHECKPOINT_ID: None} ), next_checkpoint, { "source": "fork", "step": step + 1, "parents": saved.metadata.get("parents", {}), }, {}, ) # we want to both clone a checkpoint and update state in one go. # reuse the same task ID if possible. if isinstance(values, list) and len(values) > 0: # figure out the task IDs for the next update checkpoint next_tasks = prepare_next_tasks( next_checkpoint, saved.pending_writes or [], self.nodes, channels, managed, next_config, step + 2, step + 4, for_execution=True, store=self.store, checkpointer=checkpointer, manager=None, ) tasks_group_by = defaultdict(list) user_group_by: dict[str, list[StateUpdate]] = defaultdict(list) for task in next_tasks.values(): tasks_group_by[task.name].append(task.id) for item in values: if not isinstance(item, Sequence): raise InvalidUpdateError( f"Invalid update item: {item} when copying checkpoint" ) values, as_node = item[:2] user_group = user_group_by[as_node] tasks_group = tasks_group_by[as_node] target_idx = len(user_group) task_id = ( tasks_group[target_idx] if target_idx < len(tasks_group) else None ) user_group_by[as_node].append( StateUpdate(values=values, as_node=as_node, task_id=task_id) ) return await aperform_superstep( patch_checkpoint_map(next_config, saved.metadata), [item for lst in user_group_by.values() for item in lst], ) return patch_checkpoint_map( next_config, saved.metadata if saved else None ) # task ids can be provided in the StateUpdate, but if not, # we use the task id generated by prepare_next_tasks node_to_task_ids: dict[str, deque[str]] = defaultdict(deque) if saved is not None and saved.pending_writes is not None: # we call prepare_next_tasks to discover the task IDs that # would have been generated, so we can reuse them and # properly populate task.result in state history next_tasks = prepare_next_tasks( checkpoint, saved.pending_writes, self.nodes, channels, managed, saved.config, step + 1, step + 3, for_execution=True, store=self.store, checkpointer=checkpointer, manager=None, ) # collect task ids to reuse so we can properly attach task results for t in next_tasks.values(): node_to_task_ids[t.name].append(t.id) valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = [] if len(updates) == 1: values, as_node, task_id = updates[0] # find last node that updated the state, if not provided if as_node is None and len(self.nodes) == 1: as_node = tuple(self.nodes)[0] elif as_node is None and not saved: if ( isinstance(self.input_channels, str) and self.input_channels in self.nodes ): as_node = self.input_channels elif as_node is None: last_seen_by_node = sorted( (v, n) for n, seen in checkpoint["versions_seen"].items() if n in self.nodes for v in seen.values() ) # if two nodes updated the state at the same time, it's ambiguous if last_seen_by_node: if len(last_seen_by_node) == 1: as_node = last_seen_by_node[0][1] elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: as_node = last_seen_by_node[-1][1] if as_node is None: raise InvalidUpdateError("Ambiguous update, specify as_node") if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") valid_updates.append((as_node, values, task_id)) else: for values, as_node, task_id in updates: if as_node is None: raise InvalidUpdateError( "as_node is required when applying multiple updates" ) if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") valid_updates.append((as_node, values, task_id)) run_tasks: list[PregelTaskWrites] = [] run_task_ids: list[str] = [] for as_node, values, provided_task_id in valid_updates: # create task to run all writers of the chosen node writers = self.nodes[as_node].flat_writers if not writers: raise InvalidUpdateError(f"Node {as_node} has no writers") writes: deque[tuple[str, Any]] = deque() task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) # get the task ids that were prepared for this node # if a task id was provided in the StateUpdate, we use it # otherwise, we use the next available task id prepared_task_ids = node_to_task_ids.get(as_node, deque()) task_id = provided_task_id or ( prepared_task_ids.popleft() if prepared_task_ids else str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) ) run_tasks.append(task) run_task_ids.append(task_id) run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] # execute task await run.ainvoke( values, patch_config( config, run_name=self.name + "UpdateState", configurable={ # deque.extend is thread-safe CONFIG_KEY_SEND: writes.extend, CONFIG_KEY_TASK_ID: task_id, CONFIG_KEY_READ: partial( local_read, _scratchpad( None, [], task_id, "", None, step, step + 2, ), channels, managed, task, ), }, ), ) # save task writes for task_id, task in zip(run_task_ids, run_tasks): # channel writes are saved to current checkpoint channel_writes = [w for w in task.writes if w[0] != PUSH] if saved and channel_writes: await checkpointer.aput_writes( checkpoint_config, channel_writes, task_id ) # apply to checkpoint and save apply_writes( checkpoint, channels, run_tasks, checkpointer.get_next_version, self.trigger_to_nodes, ) checkpoint = create_checkpoint(checkpoint, channels, step + 1) # save checkpoint, after applying writes next_config = await checkpointer.aput( checkpoint_config, checkpoint, { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) for task_id, task in zip(run_task_ids, run_tasks): # save push writes if push_writes := [w for w in task.writes if w[0] == PUSH]: await checkpointer.aput_writes(next_config, push_writes, task_id) return patch_checkpoint_map(next_config, saved.metadata if saved else None) current_config = patch_configurable( config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])} ) for superstep in supersteps: current_config = await aperform_superstep(current_config, superstep) return current_config def update_state( self, config: RunnableConfig, values: dict[str, Any] | Any | None, as_node: str | None = None, task_id: str | None = None, ) -> RunnableConfig: """Update the state of the graph with the given values, as if they came from node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. """ return self.bulk_update_state(config, [[StateUpdate(values, as_node, task_id)]]) async def aupdate_state( self, config: RunnableConfig, values: dict[str, Any] | Any, as_node: str | None = None, task_id: str | None = None, ) -> RunnableConfig: """Asynchronously update the state of the graph with the given values, as if they came from node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. """ return await self.abulk_update_state( config, [[StateUpdate(values, as_node, task_id)]] ) def _defaults( self, config: RunnableConfig, *, stream_mode: StreamMode | Sequence[StreamMode], print_mode: StreamMode | Sequence[StreamMode], output_keys: str | Sequence[str] | None, interrupt_before: All | Sequence[str] | None, interrupt_after: All | Sequence[str] | None, durability: Durability | None = None, ) -> tuple[ set[StreamMode], str | Sequence[str], All | Sequence[str], All | Sequence[str], BaseCheckpointSaver | None, BaseStore | None, BaseCache | None, Durability, ]: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if output_keys is None: output_keys = self.stream_channels_asis else: validate_keys(output_keys, self.channels) interrupt_before = interrupt_before or self.interrupt_before_nodes interrupt_after = interrupt_after or self.interrupt_after_nodes if isinstance(stream_mode, str): stream_modes = {stream_mode} else: stream_modes = set(stream_mode) if isinstance(print_mode, str): stream_modes.add(print_mode) else: stream_modes.update(print_mode) if self.checkpointer is False: checkpointer: BaseCheckpointSaver | None = None elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}): checkpointer = config[CONF][CONFIG_KEY_CHECKPOINTER] elif self.checkpointer is True: raise RuntimeError("checkpointer=True cannot be used for root graphs.") else: checkpointer = self.checkpointer if checkpointer and not config.get(CONF): raise ValueError( "Checkpointer requires one or more of the following 'configurable' " "keys: thread_id, checkpoint_ns, checkpoint_id" ) if CONFIG_KEY_RUNTIME in config.get(CONF, {}): store: BaseStore | None = config[CONF][CONFIG_KEY_RUNTIME].store else: store = self.store if CONFIG_KEY_CACHE in config.get(CONF, {}): cache: BaseCache | None = config[CONF][CONFIG_KEY_CACHE] else: cache = self.cache if durability is None: durability = config.get(CONF, {}).get(CONFIG_KEY_DURABILITY, "async") return ( stream_modes, output_keys, interrupt_before, interrupt_after, checkpointer, store, cache, durability, ) def stream( self, input: InputT | Command | None, config: RunnableConfig | None = None, *, context: ContextT | None = None, stream_mode: StreamMode | Sequence[StreamMode] | None = None, print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> Iterator[dict[str, Any] | Any]: """Stream graph steps for a single input. Args: input: The input to the graph. config: The configuration to use for the run. context: The static context to use for the run. !!! version-added "Added in version 0.6.0" stream_mode: The mode to stream output, defaults to `self.stream_mode`. Options are: - `"values"`: Emit all values in the state after each step, including interrupts. When used with functional API, values are emitted once at the end of the workflow. - `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step. If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately. - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`. - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks. - Will be emitted as 2-tuples `(LLM token, metadata)`. - `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`. - `"tasks"`: Emit events when tasks start and finish, including their results and errors. - `"debug"`: Emit debug events with as much information as possible for each step. You can pass a list as the `stream_mode` parameter to stream multiple modes at once. The streamed outputs will be tuples of `(mode, data)`. See [LangGraph streaming guide](https://docs.langchain.com/oss/python/langgraph/streaming) for more details. print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. durability: The durability mode for the graph execution, defaults to `"async"`. Options are: - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. subgraphs: Whether to stream events from inside subgraphs, defaults to `False`. If `True`, the events will be emitted as tuples `(namespace, data)`, or `(namespace, mode, data)` if `stream_mode` is a list, where `namespace` is a tuple with the path to the node where a subgraph is invoked, e.g. `("parent_node:", "child_node:")`. See [LangGraph streaming guide](https://docs.langchain.com/oss/python/langgraph/streaming) for more details. Yields: The output of each step in the graph. The output shape depends on the `stream_mode`. """ if (checkpoint_during := kwargs.get("checkpoint_during")) is not None: warnings.warn( "`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.", category=LangGraphDeprecatedSinceV10, stacklevel=2, ) if durability is not None: raise ValueError( "Cannot use both `checkpoint_during` and `durability` parameters. Please use `durability` instead." ) durability = "async" if checkpoint_during else "exit" if stream_mode is None: # if being called as a node in another graph, default to values mode # but don't overwrite stream_mode arg if provided stream_mode = ( "values" if config is not None and CONFIG_KEY_TASK_ID in config.get(CONF, {}) else self.stream_mode ) if debug or self.debug: print_mode = ["updates", "values"] stream = SyncQueue() config = ensure_config(self.config, config) callback_manager = get_callback_manager_for_config(config) run_manager = callback_manager.on_chain_start( None, input, name=config.get("run_name", self.get_name()), run_id=config.get("run_id"), ) try: # assign defaults ( stream_modes, output_keys, interrupt_before_, interrupt_after_, checkpointer, store, cache, durability_, ) = self._defaults( config, stream_mode=stream_mode, print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, ) if checkpointer is None and durability is not None: warnings.warn( "`durability` has no effect when no checkpointer is present.", ) # set up subgraph checkpointing if self.checkpointer is True: ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) # set up messages stream mode if "messages" in stream_modes: ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) run_manager.inheritable_handlers.append( StreamMessagesHandler( stream.put, subgraphs, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, ) ) # set up custom stream mode if "custom" in stream_modes: def stream_writer(c: Any) -> None: stream.put( ( tuple( get_config()[CONF][CONFIG_KEY_CHECKPOINT_NS].split( NS_SEP )[:-1] ), "custom", c, ) ) elif CONFIG_KEY_STREAM in config[CONF]: stream_writer = config[CONF][CONFIG_KEY_RUNTIME].stream_writer else: def stream_writer(c: Any) -> None: pass # set durability mode for subgraphs if durability is not None: config[CONF][CONFIG_KEY_DURABILITY] = durability_ runtime = Runtime( context=_coerce_context(self.context_schema, context), store=store, stream_writer=stream_writer, previous=None, ) parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME) runtime = parent_runtime.merge(runtime) config[CONF][CONFIG_KEY_RUNTIME] = runtime with SyncPregelLoop( input, stream=StreamProtocol(stream.put, stream_modes), config=config, store=store, cache=cache, checkpointer=checkpointer, nodes=self.nodes, specs=self.channels, 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, durability=durability_, trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, ) as loop: # create runner runner = PregelRunner( submit=config[CONF].get( CONFIG_KEY_RUNNER_SUBMIT, weakref.WeakMethod(loop.submit) ), put_writes=weakref.WeakMethod(loop.put_writes), node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming if subgraphs: loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream # enable concurrent streaming get_waiter: Callable[[], concurrent.futures.Future[None]] | None = None if ( self.stream_eager or subgraphs or "messages" in stream_modes or "custom" in stream_modes ): # we are careful to have a single waiter live at any one time # because on exit we increment semaphore count by exactly 1 waiter: concurrent.futures.Future | None = None # because sync futures cannot be cancelled, we instead # release the stream semaphore on exit, which will cause # a pending waiter to return immediately loop.stack.callback(stream._count.release) def get_waiter() -> concurrent.futures.Future[None]: nonlocal waiter if waiter is None or waiter.done(): waiter = loop.submit(stream.wait) return waiter else: return waiter # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates. # Channel updates from step N are only visible in step N+1 # 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(): for task in loop.match_cached_writes(): loop.output_writes(task.id, task.writes, cached=True) for _ in runner.tick( [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, get_waiter=get_waiter, schedule_task=loop.accept_push, ): # emit output yield from _output( stream_mode, print_mode, subgraphs, stream.get, queue.Empty ) loop.after_tick() # wait for checkpoint if durability_ == "sync": loop._put_checkpoint_fut.result() # emit output yield from _output( stream_mode, print_mode, subgraphs, stream.get, queue.Empty ) # handle exit if loop.status == "out_of_steps": msg = create_error_message( message=( f"Recursion limit of {config['recursion_limit']} reached " "without hitting a stop condition. You can increase the " "limit by setting the `recursion_limit` config key." ), error_code=ErrorCode.GRAPH_RECURSION_LIMIT, ) raise GraphRecursionError(msg) # 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 async def astream( self, input: InputT | Command | None, config: RunnableConfig | None = None, *, context: ContextT | None = None, stream_mode: StreamMode | Sequence[StreamMode] | None = None, print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, **kwargs: Unpack[DeprecatedKwargs], ) -> AsyncIterator[dict[str, Any] | Any]: """Asynchronously stream graph steps for a single input. Args: input: The input to the graph. config: The configuration to use for the run. context: The static context to use for the run. !!! version-added "Added in version 0.6.0" stream_mode: The mode to stream output, defaults to `self.stream_mode`. Options are: - `"values"`: Emit all values in the state after each step, including interrupts. When used with functional API, values are emitted once at the end of the workflow. - `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step. If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately. - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`. - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks. - Will be emitted as 2-tuples `(LLM token, metadata)`. - `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`. - `"tasks"`: Emit events when tasks start and finish, including their results and errors. - `"debug"`: Emit debug events with as much information as possible for each step. You can pass a list as the `stream_mode` parameter to stream multiple modes at once. The streamed outputs will be tuples of `(mode, data)`. See [LangGraph streaming guide](https://docs.langchain.com/oss/python/langgraph/streaming) for more details. print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. durability: The durability mode for the graph execution, defaults to `"async"`. Options are: - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. subgraphs: Whether to stream events from inside subgraphs, defaults to `False`. If `True`, the events will be emitted as tuples `(namespace, data)`, or `(namespace, mode, data)` if `stream_mode` is a list, where `namespace` is a tuple with the path to the node where a subgraph is invoked, e.g. `("parent_node:", "child_node:")`. See [LangGraph streaming guide](https://docs.langchain.com/oss/python/langgraph/streaming) for more details. Yields: The output of each step in the graph. The output shape depends on the `stream_mode`. """ if (checkpoint_during := kwargs.get("checkpoint_during")) is not None: warnings.warn( "`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.", category=LangGraphDeprecatedSinceV10, stacklevel=2, ) if durability is not None: raise ValueError( "Cannot use both `checkpoint_during` and `durability` parameters. Please use `durability` instead." ) durability = "async" if checkpoint_during else "exit" if stream_mode is None: # if being called as a node in another graph, default to values mode # but don't overwrite stream_mode arg if provided stream_mode = ( "values" if config is not None and CONFIG_KEY_TASK_ID in config.get(CONF, {}) else self.stream_mode ) if debug or self.debug: print_mode = ["updates", "values"] stream = AsyncQueue() aioloop = asyncio.get_running_loop() stream_put = cast( Callable[[StreamChunk], None], partial(aioloop.call_soon_threadsafe, stream.put_nowait), ) config = ensure_config(self.config, config) callback_manager = get_async_callback_manager_for_config(config) run_manager = await callback_manager.on_chain_start( None, input, name=config.get("run_name", self.get_name()), run_id=config.get("run_id"), ) # if running from astream_log() run each proc with streaming do_stream = ( next( ( True for h in run_manager.handlers if isinstance(h, _StreamingCallbackHandler) and not isinstance(h, StreamMessagesHandler) ), False, ) if _StreamingCallbackHandler is not None else False ) try: # assign defaults ( stream_modes, output_keys, interrupt_before_, interrupt_after_, checkpointer, store, cache, durability_, ) = self._defaults( config, stream_mode=stream_mode, print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, ) if checkpointer is None and durability is not None: warnings.warn( "`durability` has no effect when no checkpointer is present.", ) # set up subgraph checkpointing if self.checkpointer is True: ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) # set up messages stream mode if "messages" in stream_modes: # namespace can be None in a root level graph? ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) run_manager.inheritable_handlers.append( StreamMessagesHandler( stream_put, subgraphs, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, ) ) # set up custom stream mode def stream_writer(c: Any) -> None: aioloop.call_soon_threadsafe( stream.put_nowait, ( tuple( get_config()[CONF][CONFIG_KEY_CHECKPOINT_NS].split(NS_SEP)[ :-1 ] ), "custom", c, ), ) if "custom" in stream_modes: def stream_writer(c: Any) -> None: aioloop.call_soon_threadsafe( stream.put_nowait, ( tuple( get_config()[CONF][CONFIG_KEY_CHECKPOINT_NS].split( NS_SEP )[:-1] ), "custom", c, ), ) elif CONFIG_KEY_STREAM in config[CONF]: stream_writer = config[CONF][CONFIG_KEY_RUNTIME].stream_writer else: def stream_writer(c: Any) -> None: pass # set durability mode for subgraphs if durability is not None: config[CONF][CONFIG_KEY_DURABILITY] = durability_ runtime = Runtime( context=_coerce_context(self.context_schema, context), store=store, stream_writer=stream_writer, previous=None, ) parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME) runtime = parent_runtime.merge(runtime) config[CONF][CONFIG_KEY_RUNTIME] = runtime async with AsyncPregelLoop( input, stream=StreamProtocol(stream.put_nowait, stream_modes), config=config, store=store, cache=cache, checkpointer=checkpointer, nodes=self.nodes, specs=self.channels, 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, durability=durability_, trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, ) as loop: # create runner runner = PregelRunner( submit=config[CONF].get( CONFIG_KEY_RUNNER_SUBMIT, weakref.WeakMethod(loop.submit) ), put_writes=weakref.WeakMethod(loop.put_writes), use_astream=do_stream, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming if subgraphs: loop.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol( stream_put, stream_modes ) # enable concurrent streaming get_waiter: Callable[[], asyncio.Task[None]] | None = None _cleanup_waiter: Callable[[], Awaitable[None]] | None = None if ( self.stream_eager or subgraphs or "messages" in stream_modes or "custom" in stream_modes ): # Keep a single waiter task alive; ensure cleanup on exit. waiter: asyncio.Task[None] | None = None def get_waiter() -> asyncio.Task[None]: nonlocal waiter if waiter is None or waiter.done(): waiter = aioloop.create_task(stream.wait()) def _clear(t: asyncio.Task[None]) -> None: nonlocal waiter if waiter is t: waiter = None waiter.add_done_callback(_clear) return waiter async def _cleanup_waiter() -> None: """Wake pending waiter and/or cancel+await to avoid pending tasks.""" nonlocal waiter # Try to wake via semaphore like SyncPregelLoop with contextlib.suppress(Exception): if hasattr(stream, "_count"): stream._count.release() t = waiter waiter = None if t is not None and not t.done(): t.cancel() with contextlib.suppress(asyncio.CancelledError): await t # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps try: while loop.tick(): for task in await loop.amatch_cached_writes(): loop.output_writes(task.id, task.writes, cached=True) async for _ in runner.atick( [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, get_waiter=get_waiter, schedule_task=loop.aaccept_push, ): # emit output for o in _output( stream_mode, print_mode, subgraphs, stream.get_nowait, asyncio.QueueEmpty, ): yield o loop.after_tick() # wait for checkpoint if durability_ == "sync": await cast(asyncio.Future, loop._put_checkpoint_fut) finally: # ensure waiter doesn't remain pending on cancel/shutdown if _cleanup_waiter is not None: await _cleanup_waiter() # emit output for o in _output( stream_mode, print_mode, subgraphs, stream.get_nowait, asyncio.QueueEmpty, ): yield o # handle exit if loop.status == "out_of_steps": msg = create_error_message( message=( f"Recursion limit of {config['recursion_limit']} reached " "without hitting a stop condition. You can increase the " "limit by setting the `recursion_limit` config key." ), error_code=ErrorCode.GRAPH_RECURSION_LIMIT, ) raise GraphRecursionError(msg) # set final channel values as run output await run_manager.on_chain_end(loop.output) except BaseException as e: await asyncio.shield(run_manager.on_chain_error(e)) raise def invoke( self, input: InputT | Command | None, config: RunnableConfig | None = None, *, context: ContextT | None = None, stream_mode: StreamMode = "values", print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, **kwargs: Any, ) -> dict[str, Any] | Any: """Run the graph with a single input and config. Args: input: The input data for the graph. It can be a dictionary or any other type. config: The configuration for the graph run. context: The static context to use for the run. !!! version-added "Added in version 0.6.0" stream_mode: The stream mode for the graph run. print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. output_keys: The output keys to retrieve from the graph run. interrupt_before: The nodes to interrupt the graph run before. interrupt_after: The nodes to interrupt the graph run after. durability: The durability mode for the graph execution, defaults to `"async"`. Options are: - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. **kwargs: Additional keyword arguments to pass to the graph run. Returns: The output of the graph run. If `stream_mode` is `"values"`, it returns the latest output. If `stream_mode` is not `"values"`, it returns a list of output chunks. """ output_keys = output_keys if output_keys is not None else self.output_channels latest: dict[str, Any] | Any = None chunks: list[dict[str, Any] | Any] = [] interrupts: list[Interrupt] = [] for chunk in self.stream( input, config, context=context, stream_mode=["updates", "values"] if stream_mode == "values" else stream_mode, print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, **kwargs, ): if stream_mode == "values": if len(chunk) == 2: mode, payload = cast(tuple[StreamMode, Any], chunk) else: _, mode, payload = cast( tuple[tuple[str, ...], StreamMode, Any], chunk ) if ( mode == "updates" and isinstance(payload, dict) and (ints := payload.get(INTERRUPT)) is not None ): interrupts.extend(ints) elif mode == "values": latest = payload else: chunks.append(chunk) if stream_mode == "values": if interrupts: return ( {**latest, INTERRUPT: interrupts} if isinstance(latest, dict) else {INTERRUPT: interrupts} ) return latest else: return chunks async def ainvoke( self, input: InputT | Command | None, config: RunnableConfig | None = None, *, context: ContextT | None = None, stream_mode: StreamMode = "values", print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, **kwargs: Any, ) -> dict[str, Any] | Any: """Asynchronously run the graph with a single input and config. Args: input: The input data for the graph. It can be a dictionary or any other type. config: The configuration for the graph run. context: The static context to use for the run. !!! version-added "Added in version 0.6.0" stream_mode: The stream mode for the graph run. print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. output_keys: The output keys to retrieve from the graph run. interrupt_before: The nodes to interrupt the graph run before. interrupt_after: The nodes to interrupt the graph run after. durability: The durability mode for the graph execution, defaults to `"async"`. Options are: - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. **kwargs: Additional keyword arguments to pass to the graph run. Returns: The output of the graph run. If `stream_mode` is `"values"`, it returns the latest output. If `stream_mode` is not `"values"`, it returns a list of output chunks. """ output_keys = output_keys if output_keys is not None else self.output_channels latest: dict[str, Any] | Any = None chunks: list[dict[str, Any] | Any] = [] interrupts: list[Interrupt] = [] async for chunk in self.astream( input, config, context=context, stream_mode=["updates", "values"] if stream_mode == "values" else stream_mode, print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, durability=durability, **kwargs, ): if stream_mode == "values": if len(chunk) == 2: mode, payload = cast(tuple[StreamMode, Any], chunk) else: _, mode, payload = cast( tuple[tuple[str, ...], StreamMode, Any], chunk ) if ( mode == "updates" and isinstance(payload, dict) and (ints := payload.get(INTERRUPT)) is not None ): interrupts.extend(ints) elif mode == "values": latest = payload else: chunks.append(chunk) if stream_mode == "values": if interrupts: return ( {**latest, INTERRUPT: interrupts} if isinstance(latest, dict) else {INTERRUPT: interrupts} ) return latest else: return chunks def clear_cache(self, nodes: Sequence[str] | None = None) -> None: """Clear the cache for the given nodes.""" if not self.cache: raise ValueError("No cache is set for this graph. Cannot clear cache.") nodes = nodes or self.nodes.keys() # collect namespaces to clear namespaces: list[tuple[str, ...]] = [] for node in nodes: if node in self.nodes: namespaces.append( ( CACHE_NS_WRITES, (identifier(self.nodes[node]) or "__dynamic__"), node, ), ) # clear cache self.cache.clear(namespaces) async def aclear_cache(self, nodes: Sequence[str] | None = None) -> None: """Asynchronously clear the cache for the given nodes.""" if not self.cache: raise ValueError("No cache is set for this graph. Cannot clear cache.") nodes = nodes or self.nodes.keys() # collect namespaces to clear namespaces: list[tuple[str, ...]] = [] for node in nodes: if node in self.nodes: namespaces.append( ( CACHE_NS_WRITES, (identifier(self.nodes[node]) or "__dynamic__"), node, ), ) # clear cache await self.cache.aclear(namespaces) def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: """Index from a trigger to nodes that depend on it.""" trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list) for name, node in nodes.items(): for trigger in node.triggers: trigger_to_nodes[trigger].append(name) return dict(trigger_to_nodes) def _output( stream_mode: StreamMode | Sequence[StreamMode], print_mode: StreamMode | Sequence[StreamMode], stream_subgraphs: bool, getter: Callable[[], tuple[tuple[str, ...], str, Any]], empty_exc: type[Exception], ) -> Iterator: while True: try: ns, mode, payload = getter() except empty_exc: break if mode in print_mode: if stream_subgraphs and ns: print( " ".join( ( get_bolded_text(f"[{mode}]"), get_colored_text(f"[graph={ns}]", color="yellow"), repr(payload), ) ) ) else: print( " ".join( ( get_bolded_text(f"[{mode}]"), repr(payload), ) ) ) if mode in stream_mode: if stream_subgraphs and isinstance(stream_mode, list): yield (ns, mode, payload) elif isinstance(stream_mode, list): yield (mode, payload) elif stream_subgraphs: yield (ns, payload) else: yield payload def _coerce_context( context_schema: type[ContextT] | None, context: Any ) -> ContextT | None: """Coerce context input to the appropriate schema type. If context is a dict and context_schema is a dataclass or pydantic model, we coerce. Else, we return the context as-is. Args: context_schema: The schema type to coerce to (BaseModel, dataclass, or TypedDict) context: The context value to coerce Returns: The coerced context value or None if context is None """ if context is None: return None if context_schema is None: return context schema_is_class = issubclass(context_schema, BaseModel) or is_dataclass( context_schema ) if isinstance(context, dict) and schema_is_class: return context_schema(**context) # type: ignore[misc] return cast(ContextT, context)