4315 lines
169 KiB
Python
4315 lines
169 KiB
Python
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, replace
|
||
from datetime import timedelta
|
||
from functools import partial
|
||
from inspect import isclass
|
||
from typing import (
|
||
Any,
|
||
Generic,
|
||
Literal,
|
||
cast,
|
||
get_type_hints,
|
||
overload,
|
||
)
|
||
from uuid import UUID, uuid5
|
||
|
||
from langchain_core._api import beta
|
||
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 langchain_core.runnables.schema import StreamEvent
|
||
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 import _serde
|
||
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_STREAM_MESSAGES_V2,
|
||
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._timeout import coerce_timeout_policy
|
||
from langgraph._internal._typing import MISSING, DeprecatedKwargs
|
||
from langgraph.callbacks import (
|
||
GraphInterruptEvent,
|
||
GraphResumeEvent,
|
||
get_async_graph_callback_manager_for_config,
|
||
get_sync_graph_callback_manager_for_config,
|
||
)
|
||
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,
|
||
GraphDrained,
|
||
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 (
|
||
achannels_from_checkpoint,
|
||
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,
|
||
StreamMessagesHandlerV2,
|
||
)
|
||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||
from langgraph.pregel._retry import RetryPolicy
|
||
from langgraph.pregel._runner import PregelRunner
|
||
from langgraph.pregel._tools import StreamToolCallHandler
|
||
from langgraph.pregel._utils import (
|
||
get_new_channel_versions,
|
||
validate_timeout_supported,
|
||
)
|
||
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,
|
||
BaseUser,
|
||
RunControl,
|
||
Runtime,
|
||
ServerInfo,
|
||
)
|
||
from langgraph.stream._mux import StreamMux
|
||
from langgraph.stream._types import StreamTransformer
|
||
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
|
||
from langgraph.stream.transformers import (
|
||
LifecycleTransformer,
|
||
MessagesTransformer,
|
||
SubgraphTransformer,
|
||
ValuesTransformer,
|
||
)
|
||
from langgraph.types import (
|
||
All,
|
||
CachePolicy,
|
||
Checkpointer,
|
||
Command,
|
||
Durability,
|
||
GraphOutput,
|
||
Interrupt,
|
||
Send,
|
||
StateSnapshot,
|
||
StateUpdate,
|
||
StreamMode,
|
||
StreamPart,
|
||
TimeoutPolicy,
|
||
ensure_valid_checkpointer,
|
||
)
|
||
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",
|
||
"_timeout",
|
||
)
|
||
|
||
_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
|
||
_timeout: TimeoutPolicy | 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
|
||
self._timeout = 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 set_timeout(self, timeout: float | timedelta | TimeoutPolicy | None) -> Self:
|
||
"""Set the per-attempt timeout policy for this node."""
|
||
self._timeout = coerce_timeout_policy(timeout)
|
||
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,
|
||
timeout=self._timeout,
|
||
)
|
||
|
||
|
||
# Kwargs that ``stream_events(version="v3")`` / ``astream_events(version="v3")``
|
||
# manage internally and must not be overridden by callers. ``stream_mode`` is
|
||
# derived from the transformer mux; ``subgraphs`` is forced True so nested
|
||
# namespaces flow through scoped muxes. Forwarding either to the inner
|
||
# ``stream(...)`` would silently break v3's invariants, so we raise instead.
|
||
_V3_INVARIANT_KWARGS: tuple[str, ...] = ("stream_mode", "subgraphs")
|
||
|
||
|
||
def _reject_v3_invariant_kwargs(kwargs: dict[str, Any]) -> None:
|
||
collisions = [k for k in _V3_INVARIANT_KWARGS if k in kwargs]
|
||
if collisions:
|
||
raise TypeError(
|
||
"stream_events(version='v3') / astream_events(version='v3') do "
|
||
f"not accept {', '.join(collisions)}; v3 owns these "
|
||
"(stream_mode is built from the transformer mux, subgraphs is "
|
||
"forced True so nested namespaces flow through scoped muxes)."
|
||
)
|
||
|
||
|
||
def _collect_stream_modes(mux: Any) -> list[StreamMode]:
|
||
"""Return the union of `required_stream_modes` across registered transformers.
|
||
|
||
Transformers declare the stream modes they need to function, and
|
||
`stream_events(version="v3")` asks the graph for exactly that union — no hardcoded
|
||
default set. If zero transformers declare a given mode, the graph
|
||
does not stream events for it.
|
||
"""
|
||
modes: set[StreamMode] = set()
|
||
for transformer in mux._transformers:
|
||
modes.update(
|
||
cast(
|
||
"tuple[StreamMode, ...]",
|
||
getattr(transformer, "required_stream_modes", ()),
|
||
)
|
||
)
|
||
return list(modes)
|
||
|
||
|
||
def _normalize_stream_transformer_factories(
|
||
specs: Sequence[Callable[[tuple[str, ...]], Any]] | None,
|
||
) -> list[Callable[[tuple[str, ...]], Any]]:
|
||
"""Normalize stream transformer specs to scoped factories.
|
||
|
||
A stream transformer spec is a callable that accepts
|
||
`scope: tuple[str, ...]` and returns a fresh `StreamTransformer`.
|
||
Transformer classes work when their constructor follows the same
|
||
shape. Pre-built instances are rejected because they cannot be
|
||
cloned into subgraph scopes.
|
||
"""
|
||
factories: list[Callable[[tuple[str, ...]], Any]] = []
|
||
for spec in specs or ():
|
||
if isinstance(spec, StreamTransformer):
|
||
raise TypeError(
|
||
"stream_events(version='v3') transformers must be scope-aware callables, "
|
||
f"got pre-built instance {type(spec).__name__}. Pass the "
|
||
"transformer class or a factory like "
|
||
"`lambda scope: MyTransformer(scope, ...)`."
|
||
)
|
||
if not callable(spec):
|
||
raise TypeError(
|
||
"stream_events(version='v3') transformers must be scope-aware callables, "
|
||
f"got {type(spec).__name__}."
|
||
)
|
||
|
||
def factory(scope: tuple[str, ...], _spec: Callable[..., Any] = spec) -> Any:
|
||
return _spec(scope)
|
||
|
||
factories.append(factory)
|
||
return factories
|
||
|
||
|
||
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]]
|
||
node_error_handler_map: Mapping[str, 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: Checkpointer = 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,
|
||
node_error_handler_map: Mapping[str, str] | None = None,
|
||
name: str = "LangGraph",
|
||
stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**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)
|
||
|
||
checkpointer = ensure_valid_checkpointer(checkpointer)
|
||
|
||
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.node_error_handler_map = node_error_handler_map or {}
|
||
self.name = name
|
||
self.stream_transformers: tuple[Callable[[tuple[str, ...]], Any], ...] = tuple(
|
||
stream_transformers or ()
|
||
)
|
||
self._serde_allowlist: set[tuple[str, ...]] | None = None
|
||
if auto_validate:
|
||
self.validate()
|
||
|
||
def _apply_checkpointer_allowlist(
|
||
self, checkpointer: BaseCheckpointSaver | None
|
||
) -> BaseCheckpointSaver | None:
|
||
if not _serde.STRICT_MSGPACK_ENABLED:
|
||
return checkpointer
|
||
return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist)
|
||
|
||
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:
|
||
for name, node in self.nodes.items():
|
||
if node.timeout is not None:
|
||
validate_timeout_supported(node.node or node.bound, name=name)
|
||
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
|
||
|
||
# Mappers for v2 stream coercion (pydantic/dataclass).
|
||
# Set by CompiledStateGraph; None for base Pregel.
|
||
_output_mapper: Callable[[Any], Any] | None = None
|
||
_state_mapper: Callable[[Any], Any] | None = None
|
||
|
||
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,
|
||
saver=self.checkpointer
|
||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||
else None,
|
||
config=saved.config,
|
||
)
|
||
# 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 = await achannels_from_checkpoint(
|
||
self.channels,
|
||
saved.checkpoint,
|
||
saver=self.checkpointer
|
||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||
else None,
|
||
config=saved.config,
|
||
)
|
||
# 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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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,
|
||
saver=self.checkpointer
|
||
if saved is not None
|
||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||
else None,
|
||
config=saved.config if saved is not None else None,
|
||
)
|
||
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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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 = await achannels_from_checkpoint(
|
||
self.channels,
|
||
checkpoint,
|
||
saver=self.checkpointer
|
||
if saved is not None
|
||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||
else None,
|
||
config=saved.config if saved is not None else None,
|
||
)
|
||
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 isinstance(checkpointer, BaseCheckpointSaver):
|
||
checkpointer = self._apply_checkpointer_allowlist(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,
|
||
)
|
||
|
||
@overload
|
||
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,
|
||
control: RunControl | None = None,
|
||
subgraphs: bool = False,
|
||
debug: bool | None = None,
|
||
version: Literal["v2"],
|
||
**kwargs: Unpack[DeprecatedKwargs],
|
||
) -> Iterator[StreamPart[StateT, OutputT]]: ...
|
||
|
||
@overload
|
||
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,
|
||
control: RunControl | None = None,
|
||
subgraphs: bool = False,
|
||
debug: bool | None = None,
|
||
version: Literal["v1"] = ...,
|
||
**kwargs: Unpack[DeprecatedKwargs],
|
||
) -> Iterator[dict[str, Any] | Any]: ...
|
||
|
||
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,
|
||
control: RunControl | None = None,
|
||
subgraphs: bool = False,
|
||
debug: bool | None = None,
|
||
version: Literal["v1", "v2"] = "v1",
|
||
**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.
|
||
control: Optional run control used to request cooperative drain.
|
||
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:<task_id>", "child_node:<task_id>")`.
|
||
|
||
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)
|
||
run_manager = None
|
||
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,
|
||
)
|
||
callback_manager = get_callback_manager_for_config(config)
|
||
if "messages" in stream_modes and version != "v2":
|
||
# Strip any inherited v2 messages handler so a v1 stream
|
||
# does not get routed through the content-block event
|
||
# protocol. Leave v1 handlers in place — an outer
|
||
# stream(stream_mode="messages", subgraphs=True) relies
|
||
# on its inheritable handler to observe events emitted
|
||
# by inner stream(stream_mode="messages") calls.
|
||
callback_manager.handlers = [
|
||
h
|
||
for h in callback_manager.handlers
|
||
if not isinstance(h, StreamMessagesHandlerV2)
|
||
]
|
||
callback_manager.inheritable_handlers = [
|
||
h
|
||
for h in callback_manager.inheritable_handlers
|
||
if not isinstance(h, StreamMessagesHandlerV2)
|
||
]
|
||
if "ls_integration" not in callback_manager.metadata:
|
||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||
run_manager = callback_manager.on_chain_start(
|
||
None,
|
||
input,
|
||
name=config.get("run_name", self.get_name()),
|
||
run_id=config.get("run_id"),
|
||
)
|
||
graph_callback_manager = get_sync_graph_callback_manager_for_config(
|
||
config,
|
||
run_id=run_manager.run_id,
|
||
)
|
||
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))
|
||
use_stream_messages_v2 = bool(
|
||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||
)
|
||
messages_handler_cls = (
|
||
StreamMessagesHandlerV2
|
||
if use_stream_messages_v2
|
||
else StreamMessagesHandler
|
||
)
|
||
run_manager.inheritable_handlers.append(
|
||
messages_handler_cls(
|
||
stream.put,
|
||
subgraphs,
|
||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||
)
|
||
)
|
||
|
||
# set up tools stream mode
|
||
if "tools" in stream_modes:
|
||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||
run_manager.inheritable_handlers.append(
|
||
StreamToolCallHandler(
|
||
stream.put,
|
||
subgraphs,
|
||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools 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_
|
||
|
||
# build server_info from metadata + parent runtime
|
||
parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
|
||
server_info = _build_server_info(config, parent_runtime)
|
||
|
||
runtime = Runtime(
|
||
context=_coerce_context(self.context_schema, context),
|
||
store=store,
|
||
stream_writer=stream_writer,
|
||
previous=None,
|
||
execution_info=None,
|
||
server_info=server_info,
|
||
control=control or parent_runtime.control or RunControl(),
|
||
)
|
||
runtime = parent_runtime.merge(runtime)
|
||
config[CONF][CONFIG_KEY_RUNTIME] = runtime
|
||
|
||
# resolve mappers for v2 stream coercion
|
||
_output_mapper = self._output_mapper if version == "v2" else None
|
||
_state_mapper = self._state_mapper if version == "v2" else None
|
||
|
||
def emit_graph_lifecycle_events(loop: SyncPregelLoop) -> None:
|
||
while (event := loop._pop_lifecycle_event()) is not None:
|
||
if isinstance(event, GraphResumeEvent):
|
||
graph_callback_manager.on_resume(
|
||
replace(event, run_id=graph_callback_manager.run_id)
|
||
)
|
||
else:
|
||
graph_callback_manager.on_interrupt(
|
||
replace(event, run_id=graph_callback_manager.run_id)
|
||
)
|
||
|
||
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,
|
||
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
|
||
) as loop:
|
||
emit_graph_lifecycle_events(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),
|
||
node_error_handler_map=self.node_error_handler_map,
|
||
schedule_error_handler=loop.schedule_error_handler,
|
||
)
|
||
# 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,
|
||
version,
|
||
_output_mapper,
|
||
_state_mapper,
|
||
)
|
||
loop.after_tick()
|
||
emit_graph_lifecycle_events(loop)
|
||
# wait for checkpoint
|
||
if durability_ == "sync":
|
||
loop._put_checkpoint_fut.result()
|
||
emit_graph_lifecycle_events(loop)
|
||
# emit output
|
||
yield from _output(
|
||
stream_mode,
|
||
print_mode,
|
||
subgraphs,
|
||
stream.get,
|
||
queue.Empty,
|
||
version,
|
||
_output_mapper,
|
||
_state_mapper,
|
||
)
|
||
# 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)
|
||
elif loop.status == "draining":
|
||
if loop.control is None:
|
||
raise RuntimeError("Draining status requires run control")
|
||
raise GraphDrained(loop.control.drain_reason or "shutdown")
|
||
# set final channel values as run output
|
||
run_manager.on_chain_end(loop.output)
|
||
except BaseException as e:
|
||
if run_manager is not None:
|
||
run_manager.on_chain_error(e)
|
||
raise
|
||
|
||
@overload
|
||
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,
|
||
control: RunControl | None = None,
|
||
subgraphs: bool = False,
|
||
debug: bool | None = None,
|
||
version: Literal["v2"],
|
||
**kwargs: Unpack[DeprecatedKwargs],
|
||
) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
|
||
|
||
@overload
|
||
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,
|
||
control: RunControl | None = None,
|
||
subgraphs: bool = False,
|
||
debug: bool | None = None,
|
||
version: Literal["v1"] = ...,
|
||
**kwargs: Unpack[DeprecatedKwargs],
|
||
) -> AsyncIterator[dict[str, Any] | Any]: ...
|
||
|
||
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,
|
||
control: RunControl | None = None,
|
||
subgraphs: bool = False,
|
||
debug: bool | None = None,
|
||
version: Literal["v1", "v2"] = "v1",
|
||
**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.
|
||
control: Optional run control used to request cooperative drain.
|
||
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:<task_id>", "child_node:<task_id>")`.
|
||
|
||
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)
|
||
run_manager = None
|
||
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,
|
||
)
|
||
callback_manager = get_async_callback_manager_for_config(config)
|
||
if "messages" in stream_modes and version != "v2":
|
||
# Strip any inherited v2 messages handler so a v1 stream
|
||
# does not get routed through the content-block event
|
||
# protocol. Leave v1 handlers in place — an outer
|
||
# astream(stream_mode="messages", subgraphs=True) relies
|
||
# on its inheritable handler to observe events emitted
|
||
# by inner astream(stream_mode="messages") calls.
|
||
callback_manager.handlers = [
|
||
h
|
||
for h in callback_manager.handlers
|
||
if not isinstance(h, StreamMessagesHandlerV2)
|
||
]
|
||
callback_manager.inheritable_handlers = [
|
||
h
|
||
for h in callback_manager.inheritable_handlers
|
||
if not isinstance(h, StreamMessagesHandlerV2)
|
||
]
|
||
if "ls_integration" not in callback_manager.metadata:
|
||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||
run_manager = await callback_manager.on_chain_start(
|
||
None,
|
||
input,
|
||
name=config.get("run_name", self.get_name()),
|
||
run_id=config.get("run_id"),
|
||
)
|
||
graph_callback_manager = get_async_graph_callback_manager_for_config(
|
||
config,
|
||
run_id=run_manager.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
|
||
)
|
||
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))
|
||
use_stream_messages_v2 = bool(
|
||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||
)
|
||
messages_handler_cls = (
|
||
StreamMessagesHandlerV2
|
||
if use_stream_messages_v2
|
||
else StreamMessagesHandler
|
||
)
|
||
run_manager.inheritable_handlers.append(
|
||
messages_handler_cls(
|
||
stream_put,
|
||
subgraphs,
|
||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||
)
|
||
)
|
||
|
||
# set up tools stream mode
|
||
if "tools" in stream_modes:
|
||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||
run_manager.inheritable_handlers.append(
|
||
StreamToolCallHandler(
|
||
stream_put,
|
||
subgraphs,
|
||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools 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_
|
||
|
||
# build server_info from metadata + parent runtime
|
||
parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
|
||
server_info = _build_server_info(config, parent_runtime)
|
||
|
||
runtime = Runtime(
|
||
context=_coerce_context(self.context_schema, context),
|
||
store=store,
|
||
stream_writer=stream_writer,
|
||
previous=None,
|
||
execution_info=None,
|
||
server_info=server_info,
|
||
control=control or parent_runtime.control or RunControl(),
|
||
)
|
||
runtime = parent_runtime.merge(runtime)
|
||
config[CONF][CONFIG_KEY_RUNTIME] = runtime
|
||
|
||
# resolve mappers for v2 stream coercion
|
||
_output_mapper = self._output_mapper if version == "v2" else None
|
||
_state_mapper = self._state_mapper if version == "v2" else None
|
||
|
||
async def aemit_graph_lifecycle_events(loop: AsyncPregelLoop) -> None:
|
||
while (event := loop._pop_lifecycle_event()) is not None:
|
||
if isinstance(event, GraphResumeEvent):
|
||
await graph_callback_manager.on_resume(
|
||
GraphResumeEvent(
|
||
run_id=graph_callback_manager.run_id,
|
||
status=event.status,
|
||
checkpoint_id=event.checkpoint_id,
|
||
checkpoint_ns=event.checkpoint_ns,
|
||
)
|
||
)
|
||
else:
|
||
await graph_callback_manager.on_interrupt(
|
||
GraphInterruptEvent(
|
||
run_id=graph_callback_manager.run_id,
|
||
status=event.status,
|
||
checkpoint_id=event.checkpoint_id,
|
||
checkpoint_ns=event.checkpoint_ns,
|
||
interrupts=event.interrupts,
|
||
)
|
||
)
|
||
|
||
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,
|
||
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
|
||
) as loop:
|
||
await aemit_graph_lifecycle_events(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),
|
||
node_error_handler_map=self.node_error_handler_map,
|
||
aschedule_error_handler=loop.aschedule_error_handler,
|
||
)
|
||
# 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,
|
||
version,
|
||
_output_mapper,
|
||
_state_mapper,
|
||
):
|
||
yield o
|
||
loop.after_tick()
|
||
await aemit_graph_lifecycle_events(loop)
|
||
# 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()
|
||
|
||
await aemit_graph_lifecycle_events(loop)
|
||
|
||
# emit output
|
||
for o in _output(
|
||
stream_mode,
|
||
print_mode,
|
||
subgraphs,
|
||
stream.get_nowait,
|
||
asyncio.QueueEmpty,
|
||
version,
|
||
_output_mapper,
|
||
_state_mapper,
|
||
):
|
||
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)
|
||
elif loop.status == "draining":
|
||
if loop.control is None:
|
||
raise RuntimeError("Draining status requires run control")
|
||
raise GraphDrained(loop.control.drain_reason or "shutdown")
|
||
# set final channel values as run output
|
||
await run_manager.on_chain_end(loop.output)
|
||
except BaseException as e:
|
||
if run_manager is not None:
|
||
await asyncio.shield(run_manager.on_chain_error(e))
|
||
raise
|
||
|
||
@beta(message="The v3 streaming protocol on Pregel is experimental.")
|
||
def _pregel_stream_v3(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
interrupt_before: All | Sequence[str] | None = None,
|
||
interrupt_after: All | Sequence[str] | None = None,
|
||
control: RunControl | None = None,
|
||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**kwargs: Any,
|
||
) -> Any:
|
||
"""Internal v3 sync streaming implementation. Public entry: stream_events(version='v3').
|
||
|
||
Extra keyword arguments are forwarded to the underlying ``stream(...)``
|
||
call. The dispatcher in ``stream_events`` rejects ``stream_mode`` and
|
||
``subgraphs`` since v3 owns them (``stream_mode`` is derived from the
|
||
transformer mux; ``subgraphs`` is always True so nested namespaces
|
||
flow through scoped muxes).
|
||
|
||
!!! warning
|
||
|
||
The v3 streaming protocol is experimental and may change.
|
||
"""
|
||
parent_ns = _resolve_parent_ns(self.config, config)
|
||
compiled_factories = _normalize_stream_transformer_factories(
|
||
self.stream_transformers
|
||
)
|
||
extra_factories = _normalize_stream_transformer_factories(transformers)
|
||
mux = StreamMux(
|
||
factories=[
|
||
ValuesTransformer,
|
||
MessagesTransformer,
|
||
LifecycleTransformer,
|
||
SubgraphTransformer,
|
||
*compiled_factories,
|
||
*extra_factories,
|
||
],
|
||
scope=parent_ns,
|
||
is_async=False,
|
||
)
|
||
graph_iter = iter(
|
||
self.stream(
|
||
input,
|
||
patch_configurable(config, {CONFIG_KEY_STREAM_MESSAGES_V2: True}),
|
||
stream_mode=_collect_stream_modes(mux),
|
||
subgraphs=True,
|
||
version="v2",
|
||
interrupt_before=interrupt_before,
|
||
interrupt_after=interrupt_after,
|
||
control=control,
|
||
**kwargs,
|
||
)
|
||
)
|
||
return GraphRunStream(graph_iter, mux)
|
||
|
||
@beta(message="The v3 streaming protocol on Pregel is experimental.")
|
||
async def _apregel_stream_v3(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
interrupt_before: All | Sequence[str] | None = None,
|
||
interrupt_after: All | Sequence[str] | None = None,
|
||
control: RunControl | None = None,
|
||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**kwargs: Any,
|
||
) -> Any:
|
||
"""Internal v3 async streaming implementation. Public entry: astream_events(version='v3').
|
||
|
||
Extra keyword arguments are forwarded to the underlying ``astream(...)``
|
||
call. The dispatcher in ``astream_events`` rejects ``stream_mode`` and
|
||
``subgraphs`` since v3 owns them (``stream_mode`` is derived from the
|
||
transformer mux; ``subgraphs`` is always True so nested namespaces
|
||
flow through scoped muxes).
|
||
|
||
!!! warning
|
||
|
||
The v3 streaming protocol is experimental and may change.
|
||
"""
|
||
parent_ns = _resolve_parent_ns(self.config, config)
|
||
compiled_factories = _normalize_stream_transformer_factories(
|
||
self.stream_transformers
|
||
)
|
||
extra_factories = _normalize_stream_transformer_factories(transformers)
|
||
mux = StreamMux(
|
||
factories=[
|
||
ValuesTransformer,
|
||
MessagesTransformer,
|
||
LifecycleTransformer,
|
||
SubgraphTransformer,
|
||
*compiled_factories,
|
||
*extra_factories,
|
||
],
|
||
scope=parent_ns,
|
||
is_async=True,
|
||
)
|
||
graph_aiter = self.astream(
|
||
input,
|
||
patch_configurable(config, {CONFIG_KEY_STREAM_MESSAGES_V2: True}),
|
||
stream_mode=_collect_stream_modes(mux),
|
||
subgraphs=True,
|
||
version="v2",
|
||
interrupt_before=interrupt_before,
|
||
interrupt_after=interrupt_after,
|
||
control=control,
|
||
**kwargs,
|
||
).__aiter__()
|
||
return AsyncGraphRunStream(graph_aiter, mux)
|
||
|
||
@overload
|
||
def stream_events(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
version: Literal["v1", "v2"] = "v2",
|
||
**kwargs: Any,
|
||
) -> Iterator[StreamEvent]: ...
|
||
|
||
@overload
|
||
def stream_events(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
version: Literal["v3"],
|
||
interrupt_before: All | Sequence[str] | None = None,
|
||
interrupt_after: All | Sequence[str] | None = None,
|
||
control: RunControl | None = None,
|
||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**kwargs: Any,
|
||
) -> Any: ...
|
||
|
||
def stream_events(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
version: Literal["v1", "v2", "v3"] = "v2",
|
||
interrupt_before: All | Sequence[str] | None = None,
|
||
interrupt_after: All | Sequence[str] | None = None,
|
||
control: RunControl | None = None,
|
||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**kwargs: Any,
|
||
) -> Any:
|
||
"""Stream events from this graph.
|
||
|
||
For `version="v1"` / `"v2"`, yields `StreamEvent` dicts (see
|
||
`Runnable.stream_events`). For `version="v3"`, returns a
|
||
`GraphRunStream` whose typed projections the caller drives by
|
||
iterating — no background thread.
|
||
|
||
!!! warning
|
||
|
||
The `version="v3"` API is experimental and may change.
|
||
|
||
Builds a `StreamMux` from the built-in transformers, this
|
||
graph's compile-time `stream_transformers`, and any additional
|
||
`transformers=` supplied at the call site. `run.output`,
|
||
`run.interrupted`, and `run.interrupts` work regardless of
|
||
which transformers are registered.
|
||
|
||
Note:
|
||
Nesting v1 `stream(stream_mode="messages")` inside a node
|
||
of a `stream_events(version="v3")` run is not fully
|
||
supported. The outer v3 messages handler reroutes
|
||
`BaseChatModel.invoke` through the v2 event protocol, so
|
||
the inner v1 handler does not see `on_llm_new_token`
|
||
chunks. The inner stream still yields a finalized message
|
||
via `on_llm_end`. Use `stream_events(version="v3")` for the
|
||
inner graph as well, or call `chat_model.stream(...)`
|
||
explicitly, to get token-level streaming.
|
||
|
||
Args:
|
||
input: Graph input.
|
||
config: Optional runnable config.
|
||
version: Streaming-event schema version. `"v3"` selects the
|
||
content-block-centric streaming protocol.
|
||
interrupt_before: Nodes to interrupt before, if any. Only
|
||
used for `version="v3"`.
|
||
interrupt_after: Nodes to interrupt after, if any. Only
|
||
used for `version="v3"`.
|
||
control: Optional run control used to request cooperative
|
||
drain. Only used for `version="v3"`.
|
||
transformers: Extra transformer classes or configured
|
||
factories appended after compile-time
|
||
`stream_transformers`. Factories are called as
|
||
`factory(scope)` so they can propagate to subgraph
|
||
scopes. Only used for `version="v3"`.
|
||
**kwargs: For `version="v1"`/`"v2"`, forwarded to
|
||
`Runnable.stream_events`. For `version="v3"`, forwarded
|
||
to the underlying `stream(...)` call (e.g. `context`,
|
||
`durability`, `output_keys`, `print_mode`, `debug`).
|
||
`stream_mode` and `subgraphs` are not accepted under
|
||
`version="v3"` and raise `TypeError` if supplied; v3
|
||
owns them.
|
||
|
||
Returns:
|
||
For `version="v3"`, a `GraphRunStream` the caller iterates
|
||
to drive the run. Otherwise an `Iterator[StreamEvent]`.
|
||
"""
|
||
if version == "v3":
|
||
_reject_v3_invariant_kwargs(kwargs)
|
||
return self._pregel_stream_v3(
|
||
input,
|
||
config,
|
||
interrupt_before=interrupt_before,
|
||
interrupt_after=interrupt_after,
|
||
control=control,
|
||
transformers=transformers,
|
||
**kwargs,
|
||
)
|
||
return super().stream_events(input, config, version=version, **kwargs)
|
||
|
||
@overload
|
||
def astream_events(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
version: Literal["v1", "v2"] = "v2",
|
||
**kwargs: Any,
|
||
) -> AsyncIterator[StreamEvent]: ...
|
||
|
||
@overload
|
||
def astream_events(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
version: Literal["v3"],
|
||
interrupt_before: All | Sequence[str] | None = None,
|
||
interrupt_after: All | Sequence[str] | None = None,
|
||
control: RunControl | None = None,
|
||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**kwargs: Any,
|
||
) -> Awaitable[Any]: ...
|
||
|
||
def astream_events(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
version: Literal["v1", "v2", "v3"] = "v2",
|
||
interrupt_before: All | Sequence[str] | None = None,
|
||
interrupt_after: All | Sequence[str] | None = None,
|
||
control: RunControl | None = None,
|
||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||
**kwargs: Any,
|
||
) -> AsyncIterator[StreamEvent] | Awaitable[Any]:
|
||
"""Async variant of `stream_events`.
|
||
|
||
For `version="v3"`, returns an `AsyncGraphRunStream` whose
|
||
projections can be awaited concurrently; each subscribed cursor
|
||
drives the pump when its buffer is empty. The same nesting
|
||
limitation as the sync path applies — see `stream_events` for
|
||
details.
|
||
|
||
!!! warning
|
||
|
||
The `version="v3"` API is experimental and may change.
|
||
|
||
See `stream_events` for full argument and return documentation.
|
||
"""
|
||
if version == "v3":
|
||
_reject_v3_invariant_kwargs(kwargs)
|
||
return self._apregel_stream_v3(
|
||
input,
|
||
config,
|
||
interrupt_before=interrupt_before,
|
||
interrupt_after=interrupt_after,
|
||
control=control,
|
||
transformers=transformers,
|
||
**kwargs,
|
||
)
|
||
return super().astream_events(input, config, version=version, **kwargs)
|
||
|
||
@overload
|
||
def invoke(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
context: ContextT | None = None,
|
||
stream_mode: Literal["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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v2"],
|
||
**kwargs: Any,
|
||
) -> GraphOutput[OutputT]: ...
|
||
|
||
@overload
|
||
def invoke(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
context: ContextT | None = None,
|
||
stream_mode: StreamMode,
|
||
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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v2"],
|
||
**kwargs: Any,
|
||
) -> list[StreamPart[StateT, OutputT]]: ...
|
||
|
||
@overload
|
||
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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v1"] = ...,
|
||
**kwargs: Any,
|
||
) -> dict[str, Any] | Any: ...
|
||
|
||
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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v1", "v2"] = "v1",
|
||
**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.
|
||
control: Optional run control used to request cooperative drain.
|
||
version: The streaming format version. `"v1"` (default) returns the
|
||
traditional format, `"v2"` returns `StreamPart` typed dicts when
|
||
`stream_mode` is not `"values"`.
|
||
**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] = []
|
||
|
||
if version == "v2":
|
||
# v2: values stream parts carry interrupts directly
|
||
for chunk in self.stream(
|
||
input,
|
||
config,
|
||
context=context,
|
||
stream_mode="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,
|
||
control=control,
|
||
version=version,
|
||
**kwargs,
|
||
):
|
||
if stream_mode == "values":
|
||
latest = chunk["data"]
|
||
if chunk_ints := chunk.get("interrupts", ()):
|
||
interrupts.extend(chunk_ints) # type: ignore[arg-type]
|
||
else:
|
||
chunks.append(chunk)
|
||
else:
|
||
# v1: collect interrupts from updates stream
|
||
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,
|
||
control=control,
|
||
**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 version == "v2":
|
||
return GraphOutput(value=latest, interrupts=tuple(interrupts))
|
||
if interrupts:
|
||
return (
|
||
{**latest, INTERRUPT: interrupts}
|
||
if isinstance(latest, dict)
|
||
else {INTERRUPT: interrupts}
|
||
)
|
||
return latest
|
||
else:
|
||
return chunks
|
||
|
||
@overload
|
||
async def ainvoke(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
context: ContextT | None = None,
|
||
stream_mode: Literal["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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v2"],
|
||
**kwargs: Any,
|
||
) -> GraphOutput[OutputT]: ...
|
||
|
||
@overload
|
||
async def ainvoke(
|
||
self,
|
||
input: InputT | Command | None,
|
||
config: RunnableConfig | None = None,
|
||
*,
|
||
context: ContextT | None = None,
|
||
stream_mode: StreamMode,
|
||
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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v2"],
|
||
**kwargs: Any,
|
||
) -> list[StreamPart[StateT, OutputT]]: ...
|
||
|
||
@overload
|
||
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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v1"] = ...,
|
||
**kwargs: Any,
|
||
) -> dict[str, Any] | Any: ...
|
||
|
||
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,
|
||
control: RunControl | None = None,
|
||
version: Literal["v1", "v2"] = "v1",
|
||
**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.
|
||
control: Optional run control used to request cooperative drain.
|
||
version: The streaming format version. `"v1"` (default) returns the
|
||
traditional format, `"v2"` returns `StreamPart` typed dicts when
|
||
`stream_mode` is not `"values"`.
|
||
**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] = []
|
||
|
||
if version == "v2":
|
||
# v2: values stream parts carry interrupts directly
|
||
async for chunk in self.astream(
|
||
input,
|
||
config,
|
||
context=context,
|
||
stream_mode="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,
|
||
control=control,
|
||
version=version,
|
||
**kwargs,
|
||
):
|
||
if stream_mode == "values":
|
||
latest = chunk["data"]
|
||
if chunk_ints := chunk.get("interrupts", ()):
|
||
interrupts.extend(chunk_ints) # type: ignore[arg-type]
|
||
else:
|
||
chunks.append(chunk)
|
||
else:
|
||
# v1: collect interrupts from updates stream
|
||
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,
|
||
control=control,
|
||
**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 version == "v2":
|
||
return GraphOutput(value=latest, interrupts=tuple(interrupts))
|
||
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],
|
||
version: Literal["v1", "v2"] = "v1",
|
||
output_mapper: Callable[[Any], Any] | None = None,
|
||
state_mapper: Callable[[Any], Any] | None = None,
|
||
) -> 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 version == "v2":
|
||
if mode == "values":
|
||
# pop __interrupt__ into typed field, coerce data
|
||
ints: tuple[Interrupt, ...] = ()
|
||
if isinstance(payload, dict):
|
||
ints = payload.pop(INTERRUPT, ())
|
||
if output_mapper:
|
||
payload = output_mapper(payload)
|
||
yield {"type": mode, "ns": ns, "data": payload, "interrupts": ints}
|
||
elif mode in ("checkpoints", "debug"):
|
||
# coerce state values in checkpoint/debug payloads
|
||
if state_mapper:
|
||
_coerce_checkpoint_values(payload, state_mapper)
|
||
yield {"type": mode, "ns": ns, "data": payload}
|
||
else:
|
||
yield {"type": mode, "ns": ns, "data": payload}
|
||
elif 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_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> None:
|
||
"""Coerce `values` dicts inside checkpoint or debug payloads in-place.
|
||
|
||
Skips the initial checkpoint (where next contains ``__start__``) because
|
||
not all channels are populated yet and coercion would fail.
|
||
"""
|
||
_START = "__start__"
|
||
# debug wrapper: {"type": "checkpoint", "payload": {"values": dict, ...}}
|
||
if (
|
||
isinstance(payload, dict)
|
||
and payload.get("type") == "checkpoint"
|
||
and isinstance(payload.get("payload"), dict)
|
||
and isinstance(payload["payload"].get("values"), dict)
|
||
and _START not in payload["payload"].get("next", ())
|
||
):
|
||
payload["payload"]["values"] = mapper(payload["payload"]["values"])
|
||
# direct checkpoint payload: {"values": dict, ...}
|
||
elif (
|
||
isinstance(payload, dict)
|
||
and isinstance(payload.get("values"), dict)
|
||
and _START not in payload.get("next", ())
|
||
):
|
||
payload["values"] = mapper(payload["values"])
|
||
|
||
|
||
def _resolve_parent_ns(
|
||
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
|
||
) -> tuple[str, ...]:
|
||
"""Return the checkpoint namespace the caller is running under.
|
||
|
||
`stream_events(version="v3")` uses this to scope its native projections
|
||
(`ValuesTransformer`, `MessagesTransformer`) to events emitted at
|
||
the run's own level. A root call resolves to `()`; a call made
|
||
from inside a node carries the outer graph's task namespace so the
|
||
projection still matches its own root-level events.
|
||
"""
|
||
merged = ensure_config(graph_config, call_config)
|
||
ns = merged.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS)
|
||
if not ns:
|
||
return ()
|
||
return tuple(ns.split(NS_SEP))
|
||
|
||
|
||
def _build_server_info(
|
||
config: RunnableConfig, parent_runtime: Runtime[Any]
|
||
) -> ServerInfo | None:
|
||
"""Build ServerInfo from config configurable.
|
||
|
||
The server puts assistant_id/graph_id in config configurable and the
|
||
authenticated user dict in configurable["langgraph_auth_user"].
|
||
"""
|
||
configurable = config.get(CONF) or {}
|
||
assistant_id = configurable.get("assistant_id")
|
||
graph_id = configurable.get("graph_id")
|
||
|
||
# Read authenticated user from configurable (set by LangGraph Server).
|
||
# We prefer isinstance(BaseUser) but fall back to hasattr("identity")
|
||
# because the server's ProxyUser provides `permissions` via __getattr__,
|
||
# which Python's runtime_checkable Protocol check doesn't see.
|
||
auth_user_data = configurable.get("langgraph_auth_user")
|
||
user: BaseUser | None = None
|
||
if auth_user_data is not None:
|
||
if isinstance(auth_user_data, BaseUser) or hasattr(auth_user_data, "identity"):
|
||
user = cast(BaseUser, auth_user_data)
|
||
|
||
if assistant_id is not None or graph_id is not None or user is not None:
|
||
return ServerInfo(
|
||
assistant_id=str(assistant_id) if assistant_id else "",
|
||
graph_id=str(graph_id) if graph_id else "",
|
||
user=user,
|
||
)
|
||
return None
|
||
|
||
|
||
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)
|