Merge pull request #1792 from langchain-ai/nc/21sep/docstrings-comments

Add more comments and docstrings

docs/docs/reference/graphs.md

@@ -29,7 +29,7 @@ handler: python
 
 ## StreamMode
 
-::: langgraph.pregel.StreamMode
+::: langgraph.types.StreamMode
 
 ## Constants
 
@@ -69,8 +69,12 @@ builder.add_conditional_edges("my_node", my_condition)
 
 ## Send
 
-::: langgraph.constants.Send
+::: langgraph.types.Send
+
+## Interrupt
+
+::: langgraph.types.Interrupt
 
 ## RetryPolicy
 
-::: langgraph.pregel.types.RetryPolicy
+::: langgraph.types.RetryPolicy

libs/langgraph/langgraph/constants.py

@@ -1,133 +1,109 @@
-from dataclasses import dataclass
 from types import MappingProxyType
-from typing import Any, Literal, Mapping
+from typing import Any, Mapping
 
+from langgraph.types import Interrupt, Send  # noqa: F401
+
+# Interrupt, Send re-exported for backwards compatibility
+
+
+# --- Empty read-only containers ---
+EMPTY_MAP: Mapping[str, Any] = MappingProxyType({})
+EMPTY_SEQ: tuple[str, ...] = tuple()
+
+# --- Public constants ---
+TAG_HIDDEN = "langsmith:hidden"
+# tag to hide a node/edge from certain tracing/streaming environments
+START = "__start__"
+# the first (maybe virtual) node in graph-style Pregel
+END = "__end__"
+# the last (maybe virtual) node in graph-style Pregel
+
+# --- Reserved write keys ---
 INPUT = "__input__"
-CONFIG_KEY_SEND = "__pregel_send"
-CONFIG_KEY_READ = "__pregel_read"
-CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
-CONFIG_KEY_STREAM = "__pregel_stream"
-CONFIG_KEY_STREAM_WRITER = "__pregel_stream_writer"
-CONFIG_KEY_STORE = "__pregel_store"
-CONFIG_KEY_RESUMING = "__pregel_resuming"
-CONFIG_KEY_TASK_ID = "__pregel_task_id"
-CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks"
-CONFIG_KEY_ENSURE_LATEST = "__pregel_ensure_latest"
-CONFIG_KEY_DELEGATE = "__pregel_delegate"
-# this one part of public API so more readable
-CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map"
+# for values passed as input to the graph
 INTERRUPT = "__interrupt__"
+# for dynamic interrupts raised by nodes
 ERROR = "__error__"
+# for errors raised by nodes
 NO_WRITES = "__no_writes__"
+# marker to signal node didn't write anything
 SCHEDULED = "__scheduled__"
-TASKS = "__pregel_tasks"  # for backwards compat, this is the original name of PUSH
+# marker to signal node was scheduled (in distributed mode)
+TASKS = "__pregel_tasks"
+# for Send objects returned by nodes/edges, corresponds to PUSH below
+
+# --- Reserved config.configurable keys ---
+CONFIG_KEY_SEND = "__pregel_send"
+# holds the `write` function that accepts writes to state/edges/reserved keys
+CONFIG_KEY_READ = "__pregel_read"
+# holds the `read` function that returns a copy of the current state
+CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
+# holds a `BaseCheckpointSaver` passed from parent graph to child graphs
+CONFIG_KEY_STREAM = "__pregel_stream"
+# holds a `StreamProtocol` passed from parent graph to child graphs
+CONFIG_KEY_STREAM_WRITER = "__pregel_stream_writer"
+# holds a `StreamWriter` for stream_mode=custom
+CONFIG_KEY_STORE = "__pregel_store"
+# holds a `BaseStore` made available to managed values
+CONFIG_KEY_RESUMING = "__pregel_resuming"
+# holds a boolean indicating if subgraphs should resume from a previous checkpoint
+CONFIG_KEY_TASK_ID = "__pregel_task_id"
+# holds the task ID for the current task
+CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks"
+# holds a boolean indicating if tasks should be deduplicated (for distributed mode)
+CONFIG_KEY_ENSURE_LATEST = "__pregel_ensure_latest"
+# holds a boolean indicating whether to assert the requested checkpoint is the latest
+# (for distributed mode)
+CONFIG_KEY_DELEGATE = "__pregel_delegate"
+# holds a boolean indicating whether to delegate subgraphs (for distributed mode)
+CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map"
+# holds a mapping of checkpoint_ns -> checkpoint_id for parent graphs
+CONFIG_KEY_CHECKPOINT_ID = "checkpoint_id"
+# holds the current checkpoint_id, if any
+CONFIG_KEY_CHECKPOINT_NS = "checkpoint_ns"
+# holds the current checkpoint_ns, "" for root graph
+
+# --- Other constants ---
 PUSH = "__pregel_push"
+# denotes push-style tasks, ie. those created by Send objects
 PULL = "__pregel_pull"
+# denotes pull-style tasks, ie. those triggered by edges
 RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__"
+# placeholder for managed values replaced at runtime
+NS_SEP = "|"
+# for checkpoint_ns, separates each level (ie. graph|subgraph|subsubgraph)
+NS_END = ":"
+# for checkpoint_ns, for each level, separates the namespace from the task_id
+
 RESERVED = {
-    SCHEDULED,
+    TAG_HIDDEN,
+    # reserved write keys
+    INPUT,
     INTERRUPT,
     ERROR,
     NO_WRITES,
+    SCHEDULED,
     TASKS,
-    PUSH,
-    PULL,
+    # reserved config.configurable keys
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
     CONFIG_KEY_CHECKPOINTER,
-    CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_STREAM,
     CONFIG_KEY_STREAM_WRITER,
     CONFIG_KEY_STORE,
+    CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     CONFIG_KEY_DEDUPE_TASKS,
     CONFIG_KEY_ENSURE_LATEST,
     CONFIG_KEY_DELEGATE,
-    INPUT,
+    CONFIG_KEY_CHECKPOINT_MAP,
+    CONFIG_KEY_CHECKPOINT_ID,
+    CONFIG_KEY_CHECKPOINT_NS,
+    # other constants
+    PUSH,
+    PULL,
     RUNTIME_PLACEHOLDER,
+    NS_SEP,
+    NS_END,
 }
-TAG_HIDDEN = "langsmith:hidden"
-
-START = "__start__"
-END = "__end__"
-
-NS_SEP = "|"
-NS_END = ":"
-
-EMPTY_MAP: Mapping[str, Any] = MappingProxyType({})
-
-
-class Send:
-    """A message or packet to send to a specific node in the graph.
-
-    The `Send` class is used within a `StateGraph`'s conditional edges to
-    dynamically invoke a node with a custom state at the next step.
-
-    Importantly, the sent state can differ from the core graph's state,
-    allowing for flexible and dynamic workflow management.
-
-    One such example is a "map-reduce" workflow where your graph invokes
-    the same node multiple times in parallel with different states,
-    before aggregating the results back into the main graph's state.
-
-    Attributes:
-        node (str): The name of the target node to send the message to.
-        arg (Any): The state or message to send to the target node.
-
-    Examples:
-        >>> from typing import Annotated
-        >>> import operator
-        >>> class OverallState(TypedDict):
-        ...     subjects: list[str]
-        ...     jokes: Annotated[list[str], operator.add]
-        ...
-        >>> from langgraph.constants import Send
-        >>> from langgraph.graph import END, START
-        >>> def continue_to_jokes(state: OverallState):
-        ...     return [Send("generate_joke", {"subject": s}) for s in state['subjects']]
-        ...
-        >>> from langgraph.graph import StateGraph
-        >>> builder = StateGraph(OverallState)
-        >>> builder.add_node("generate_joke", lambda state: {"jokes": [f"Joke about {state['subject']}"]})
-        >>> builder.add_conditional_edges(START, continue_to_jokes)
-        >>> builder.add_edge("generate_joke", END)
-        >>> graph = builder.compile()
-        >>>
-        >>> # Invoking with two subjects results in a generated joke for each
-        >>> graph.invoke({"subjects": ["cats", "dogs"]})
-        {'subjects': ['cats', 'dogs'], 'jokes': ['Joke about cats', 'Joke about dogs']}
-    """
-
-    node: str
-    arg: Any
-
-    def __init__(self, /, node: str, arg: Any) -> None:
-        """
-        Initialize a new instance of the Send class.
-
-        Args:
-            node (str): The name of the target node to send the message to.
-            arg (Any): The state or message to send to the target node.
-        """
-        self.node = node
-        self.arg = arg
-
-    def __hash__(self) -> int:
-        return hash((self.node, self.arg))
-
-    def __repr__(self) -> str:
-        return f"Send(node={self.node!r}, arg={self.arg!r})"
-
-    def __eq__(self, value: object) -> bool:
-        return (
-            isinstance(value, Send)
-            and self.node == value.node
-            and self.arg == value.arg
-        )
-
-
-@dataclass
-class Interrupt:
-    value: Any
-    when: Literal["during"] = "during"

libs/langgraph/langgraph/errors.py

@@ -1,7 +1,9 @@
 from typing import Any, Sequence
 
-from langgraph.checkpoint.base import EmptyChannelError
-from langgraph.constants import Interrupt
+from langgraph.checkpoint.base import EmptyChannelError  # noqa: F401
+from langgraph.types import Interrupt
+
+# EmptyChannelError re-exported for backwards compatibility
 
 
 class GraphRecursionError(RecursionError):
@@ -24,13 +26,14 @@ class GraphRecursionError(RecursionError):
 
 
 class InvalidUpdateError(Exception):
-    """Raised when attempting to update a channel with an invalid sequence of updates."""
+    """Raised when attempting to update a channel with an invalid set of updates."""
 
     pass
 
 
 class GraphInterrupt(Exception):
-    """Raised when a subgraph is interrupted."""
+    """Raised when a subgraph is interrupted, suppressed by the root graph.
+    Never raised directly, or surfaced to the user."""
 
     def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None:
         super().__init__(interrupts)
@@ -44,7 +47,7 @@ class NodeInterrupt(GraphInterrupt):
 
 
 class GraphDelegate(Exception):
-    """Raised when a graph is delegated."""
+    """Raised when a graph is delegated (for distributed mode)."""
 
     def __init__(self, *args: dict[str, Any]) -> None:
         super().__init__(*args)
@@ -57,22 +60,22 @@ class EmptyInputError(Exception):
 
 
 class TaskNotFound(Exception):
-    """Raised when the executor is unable to find a task."""
+    """Raised when the executor is unable to find a task (for distributed mode)."""
 
     pass
 
 
 class CheckpointNotLatest(Exception):
-    """Raised when the checkpoint is not the latest version."""
+    """Raised when the checkpoint is not the latest version (for distributed mode)."""
 
     pass
 
 
-__all__ = [
-    "GraphRecursionError",
-    "InvalidUpdateError",
-    "GraphInterrupt",
-    "NodeInterrupt",
-    "EmptyInputError",
-    "EmptyChannelError",
-]
+class MultipleSubgraphsError(Exception):
+    """Raised when multiple subgraphs are called inside the same node."""
+
+    pass
+
+
+_SEEN_CHECKPOINT_NS: set[str] = set()
+"""Used for subgraph detection."""

libs/langgraph/langgraph/graph/graph.py

@@ -26,7 +26,6 @@ from langchain_core.runnables.graph import Node as DrawableNode
 from typing_extensions import Self
 
 from langgraph.channels.ephemeral_value import EphemeralValue
-from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import (
     END,
     NS_END,
@@ -38,8 +37,8 @@ from langgraph.constants import (
 from langgraph.errors import InvalidUpdateError
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
-from langgraph.pregel.types import All
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.types import All, Checkpointer
 from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
@@ -406,7 +405,7 @@ class Graph:
 
     def compile(
         self,
-        checkpointer: Optional[BaseCheckpointSaver] = None,
+        checkpointer: Checkpointer = None,
         interrupt_before: Optional[Union[All, list[str]]] = None,
         interrupt_after: Optional[Union[All, list[str]]] = None,
         debug: bool = False,

libs/langgraph/langgraph/graph/state.py

@@ -32,7 +32,6 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.named_barrier_value import NamedBarrierValue
-from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN
 from langgraph.errors import InvalidUpdateError
 from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send
@@ -45,9 +44,9 @@ from langgraph.managed.base import (
     is_writable_managed_value,
 )
 from langgraph.pregel.read import ChannelRead, PregelNode
-from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
+from langgraph.types import All, Checkpointer, RetryPolicy
 from langgraph.utils.fields import get_field_default
 from langgraph.utils.pydantic import create_model
 from langgraph.utils.runnable import coerce_to_runnable
@@ -400,7 +399,7 @@ class StateGraph(Graph):
 
     def compile(
         self,
-        checkpointer: Optional[BaseCheckpointSaver] = None,
+        checkpointer: Checkpointer = None,
         *,
         store: Optional[BaseStore] = None,
         interrupt_before: Optional[Union[All, list[str]]] = None,
@@ -413,7 +412,7 @@ class StateGraph(Graph):
         streamed, batched, and run asynchronously.
 
         Args:
-            checkpointer (Optional[BaseCheckpointSaver]): An optional checkpoint saver object.
+            checkpointer (Checkpointer): An optional checkpoint saver object.
                 This serves as a fully versioned "memory" for the graph, allowing
                 the graph to be paused and resumed, and replayed from any point.
             interrupt_before (Optional[Sequence[str]]): An optional list of node names to interrupt before.

libs/langgraph/langgraph/prebuilt/chat_agent_executor.py

@@ -16,13 +16,13 @@ from langchain_core.runnables import Runnable, RunnableConfig, RunnableLambda
 from langchain_core.tools import BaseTool
 
 from langgraph._api.deprecation import deprecated_parameter
-from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph import StateGraph
 from langgraph.graph.graph import CompiledGraph
 from langgraph.graph.message import add_messages
 from langgraph.managed import IsLastStep
 from langgraph.prebuilt.tool_executor import ToolExecutor
 from langgraph.prebuilt.tool_node import ToolNode
+from langgraph.types import Checkpointer
 
 
 # We create the AgentState that we will pass around
@@ -132,7 +132,7 @@ def create_react_agent(
     state_schema: Optional[StateSchemaType] = None,
     messages_modifier: Optional[MessagesModifier] = None,
     state_modifier: Optional[StateModifier] = None,
-    checkpointer: Optional[BaseCheckpointSaver] = None,
+    checkpointer: Checkpointer = None,
     interrupt_before: Optional[list[str]] = None,
     interrupt_after: Optional[list[str]] = None,
     debug: bool = False,

libs/langgraph/langgraph/pregel/__init__.py

@@ -83,11 +83,11 @@ from langgraph.pregel.messages import StreamMessagesHandler
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.runner import PregelRunner
-from langgraph.pregel.types import All, StateSnapshot, StreamMode
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
+from langgraph.types import All, Checkpointer, StateSnapshot, StreamMode
 from langgraph.utils.config import (
     ensure_config,
     merge_configs,
@@ -197,7 +197,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
     debug: bool
     """Whether to print debug information during execution. Defaults to False."""
 
-    checkpointer: Optional[BaseCheckpointSaver] = None
+    checkpointer: Checkpointer = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
     store: Optional[BaseStore] = None
@@ -281,7 +281,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
                 [spec for node in self.nodes.values() for spec in node.config_specs]
                 + (
                     self.checkpointer.config_specs
-                    if self.checkpointer is not None
+                    if isinstance(self.checkpointer, BaseCheckpointSaver)
                     else []
                 )
                 + (
@@ -1059,6 +1059,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
         Union[All, Sequence[str]],
         Optional[BaseCheckpointSaver],
     ]:
+        if config["recursion_limit"] < 1:
+            raise ValueError("recursion_limit must be at least 1")
         debug = debug if debug is not None else self.debug
         if output_keys is None:
             output_keys = self.stream_channels_asis
@@ -1072,12 +1074,16 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
         if CONFIG_KEY_TASK_ID in config.get("configurable", {}):
             # if being called as a node in another graph, always use values mode
             stream_mode = ["values"]
-        if CONFIG_KEY_CHECKPOINTER in config.get("configurable", {}):
-            checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][
-                CONFIG_KEY_CHECKPOINTER
-            ]
+        if self.checkpointer is False:
+            checkpointer: Optional[BaseCheckpointSaver] = None
+        elif CONFIG_KEY_CHECKPOINTER in config.get("configurable", {}):
+            checkpointer = config["configurable"][CONFIG_KEY_CHECKPOINTER]
         else:
             checkpointer = self.checkpointer
+        if checkpointer and not config.get("configurable"):
+            raise ValueError(
+                f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in checkpointer.config_specs]}"
+            )
         return (
             debug,
             set(stream_mode),
@@ -1193,12 +1199,6 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
             run_id=config.get("run_id"),
         )
         try:
-            if config["recursion_limit"] < 1:
-                raise ValueError("recursion_limit must be at least 1")
-            if self.checkpointer and not config.get("configurable"):
-                raise ValueError(
-                    f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in self.checkpointer.config_specs]}"
-                )
             # assign defaults
             (
                 debug,
@@ -1414,12 +1414,6 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
             None,
         )
         try:
-            if config["recursion_limit"] < 1:
-                raise ValueError("recursion_limit must be at least 1")
-            if self.checkpointer and not config.get("configurable"):
-                raise ValueError(
-                    f"Checkpointer requires one or more of the following 'configurable' keys: {[s.id for s in self.checkpointer.config_specs]}"
-                )
             # assign defaults
             (
                 debug,

libs/langgraph/langgraph/pregel/algo.py

@@ -33,6 +33,7 @@ from langgraph.constants import (
     CONFIG_KEY_READ,
     CONFIG_KEY_SEND,
     CONFIG_KEY_TASK_ID,
+    EMPTY_SEQ,
     INTERRUPT,
     NO_WRITES,
     NS_END,
@@ -50,15 +51,16 @@ from langgraph.pregel.io import read_channel, read_channels
 from langgraph.pregel.log import logger
 from langgraph.pregel.manager import ChannelsManager
 from langgraph.pregel.read import PregelNode
-from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
+from langgraph.types import All, PregelExecutableTask, PregelTask
 from langgraph.utils.config import merge_configs, patch_config
 
 GetNextVersion = Callable[[Optional[V], BaseChannel], V]
 
-EMPTY_SEQ: tuple[str, ...] = tuple()
-
 
 class WritesProtocol(Protocol):
+    """Protocol for objects containing writes to be applied to checkpoint.
+    Implemented by PregelTaskWrites and PregelExecutableTask."""
+
     @property
     def name(self) -> str: ...
 
@@ -70,6 +72,9 @@ class WritesProtocol(Protocol):
 
 
 class PregelTaskWrites(NamedTuple):
+    """Simplest implementation of WritesProtocol, for usage with writes that
+    don't originate from a runnable task, eg. graph input, update_state, etc."""
+
     name: str
     writes: Sequence[tuple[str, Any]]
     triggers: Sequence[str]
@@ -80,6 +85,7 @@ def should_interrupt(
     interrupt_nodes: Union[All, Sequence[str]],
     tasks: Iterable[PregelExecutableTask],
 ) -> list[PregelExecutableTask]:
+    """Check if the graph should be interrupted based on current state."""
     version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
     null_version = version_type()  # type: ignore[misc]
     seen = checkpoint["versions_seen"].get(INTERRUPT, {})
@@ -117,6 +123,9 @@ def local_read(
     select: Union[list[str], str],
     fresh: bool = False,
 ) -> Union[dict[str, Any], Any]:
+    """Function injected under CONFIG_KEY_READ in task config, to read current state.
+    Used by conditional edges to read a copy of the state with reflecting the writes
+    from that node only."""
     if isinstance(select, str):
         managed_keys = []
         for c, _ in task.writes:
@@ -153,6 +162,8 @@ def local_write(
     managed: ManagedValueMapping,
     writes: Sequence[tuple[str, Any]],
 ) -> None:
+    """Function injected under CONFIG_KEY_SEND in task config, to write to channels.
+    Validates writes and forwards them to `commit` function."""
     for chan, value in writes:
         if chan == TASKS:
             if not isinstance(value, Send):
@@ -169,6 +180,7 @@ def local_write(
 
 
 def increment(current: Optional[int], channel: BaseChannel) -> int:
+    """Default channel versioning function, increments the current int version."""
     return current + 1 if current is not None else 1
 
 
@@ -178,6 +190,9 @@ def apply_writes(
     tasks: Iterable[WritesProtocol],
     get_next_version: Optional[GetNextVersion],
 ) -> dict[str, list[Any]]:
+    """Apply writes from a set of tasks (usually the tasks from a Pregel step)
+    to the checkpoint and channels, and return managed values writes to be applied
+    externally."""
     # update seen versions
     for task in tasks:
         checkpoint["versions_seen"].setdefault(task.name, {}).update(
@@ -297,6 +312,9 @@ def prepare_next_tasks(
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
 ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
+    """Prepare the set of tasks that will make up the next Pregel step.
+    This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
+    by edges)."""
     tasks: dict[str, Union[PregelTask, PregelExecutableTask]] = {}
     # Consume pending packets
     for idx, _ in enumerate(checkpoint["pending_sends"]):
@@ -348,6 +366,8 @@ def prepare_single_task(
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
 ) -> Union[None, PregelTask, PregelExecutableTask]:
+    """Prepares a single task for the next Pregel step, given a task path, which
+    uniquely identifies a PUSH or PULL task within the graph."""
     checkpoint_id = UUID(checkpoint["id"]).bytes
     configurable = config.get("configurable", {})
     parent_ns = configurable.get("checkpoint_ns", "")
@@ -568,6 +588,7 @@ def _proc_input(
     *,
     for_execution: bool,
 ) -> Iterator[Any]:
+    """Prepare input for a PULL task, based on the process's channels and triggers."""
     # If all trigger channels subscribed by this process are not empty
     # then invoke the process with the values of all non-empty channels
     if isinstance(proc.channels, dict):

libs/langgraph/langgraph/pregel/debug.py

@@ -22,7 +22,7 @@ from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite
 from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
 from langgraph.pregel.io import read_channels
-from langgraph.pregel.types import PregelExecutableTask, PregelTask, StateSnapshot
+from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot
 
 
 class TaskPayload(TypedDict):
@@ -84,6 +84,7 @@ TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
 def map_debug_tasks(
     step: int, tasks: Iterable[PregelExecutableTask]
 ) -> Iterator[DebugOutputTask]:
+    """Produce "task" events for stream_mode=debug."""
     ts = datetime.now(timezone.utc).isoformat()
     for task in tasks:
         if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):
@@ -107,6 +108,7 @@ def map_debug_task_results(
     task_tup: tuple[PregelExecutableTask, Sequence[tuple[str, Any]]],
     stream_keys: Union[str, Sequence[str]],
 ) -> Iterator[DebugOutputTaskResult]:
+    """Produce "task_result" events for stream_mode=debug."""
     stream_channels_list = (
         [stream_keys] if isinstance(stream_keys, str) else stream_keys
     )
@@ -135,6 +137,7 @@ def map_debug_checkpoint(
     tasks: Iterable[PregelExecutableTask],
     pending_writes: list[PendingWrite],
 ) -> Iterator[DebugOutputCheckpoint]:
+    """Produce "checkpoint" events for stream_mode=debug."""
     yield {
         "type": "checkpoint",
         "timestamp": checkpoint["ts"],
@@ -213,6 +216,7 @@ def tasks_w_writes(
     pending_writes: Optional[list[PendingWrite]],
     states: Optional[dict[str, Union[RunnableConfig, StateSnapshot]]],
 ) -> tuple[PregelTask, ...]:
+    """Apply writes / subgraph states to tasks to be returned in a StateSnapshot."""
     pending_writes = pending_writes or []
     return tuple(
         PregelTask(

libs/langgraph/langgraph/pregel/executor.py

@@ -39,6 +39,13 @@ class Submit(Protocol[P, T]):
 
 
 class BackgroundExecutor(ContextManager):
+    """A context manager that runs sync tasks in the background.
+    Uses a thread pool executor to delegate tasks to separate threads.
+    On exit,
+    - cancels any (not yet started) tasks with `__cancel_on_exit__=True`
+    - waits for all tasks to finish
+    - re-raises the first exception from tasks with `__reraise_on_exit__=True`"""
+
     def __init__(self, config: RunnableConfig) -> None:
         self.stack = ExitStack()
         self.executor = self.stack.enter_context(get_executor_for_config(config))
@@ -49,7 +56,7 @@ class BackgroundExecutor(ContextManager):
         fn: Callable[P, T],
         *args: P.args,
         __name__: Optional[str] = None,  # currently not used in sync version
-        __cancel_on_exit__: bool = False,
+        __cancel_on_exit__: bool = False,  # for sync, can cancel only if not started
         __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]:
@@ -101,6 +108,14 @@ class BackgroundExecutor(ContextManager):
 
 
 class AsyncBackgroundExecutor(AsyncContextManager):
+    """A context manager that runs async tasks in the background.
+    Uses the current event loop to delegate tasks to asyncio tasks.
+    On exit,
+    - cancels any tasks with `__cancel_on_exit__=True`
+    - waits for all tasks to finish
+    - re-raises the first exception from tasks with `__reraise_on_exit__=True`
+      ignoring CancelledError"""
+
     def __init__(self) -> None:
         self.context_not_supported = sys.version_info < (3, 11)
         self.tasks: dict[asyncio.Task, tuple[bool, bool]] = {}

libs/langgraph/langgraph/pregel/io.py

@@ -3,9 +3,9 @@ from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar,
 from langchain_core.runnables.utils import AddableDict
 
 from langgraph.channels.base import BaseChannel, EmptyChannelError
-from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
+from langgraph.constants import EMPTY_SEQ, ERROR, INTERRUPT, TAG_HIDDEN
 from langgraph.pregel.log import logger
-from langgraph.pregel.types import PregelExecutableTask
+from langgraph.types import PregelExecutableTask
 
 
 def read_channel(
@@ -97,9 +97,6 @@ class AddableUpdatesDict(AddableDict):
         raise TypeError("AddableUpdatesDict does not support right-side addition")
 
 
-EMPTY_SEQ: tuple[str, ...] = tuple()
-
-
 def map_output_updates(
     output_channels: Union[str, Sequence[str]],
     tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],

libs/langgraph/langgraph/pregel/loop.py

@@ -44,6 +44,7 @@ from langgraph.constants import (
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_STREAM,
     CONFIG_KEY_TASK_ID,
+    EMPTY_SEQ,
     ERROR,
     INPUT,
     INTERRUPT,
@@ -53,10 +54,12 @@ from langgraph.constants import (
     TASKS,
 )
 from langgraph.errors import (
+    _SEEN_CHECKPOINT_NS,
     CheckpointNotLatest,
     EmptyInputError,
     GraphDelegate,
     GraphInterrupt,
+    MultipleSubgraphsError,
 )
 from langgraph.managed.base import (
     ManagedValueMapping,
@@ -93,21 +96,20 @@ from langgraph.pregel.io import (
 )
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.read import PregelNode
-from langgraph.pregel.types import All, PregelExecutableTask, StreamMode
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.store.base import BaseStore
 from langgraph.store.batch import AsyncBatchedStore
+from langgraph.types import All, PregelExecutableTask, StreamMode
 from langgraph.utils.config import patch_configurable
 
 V = TypeVar("V")
 P = ParamSpec("P")
+StreamChunk = tuple[tuple[str, ...], str, Any]
+
 INPUT_DONE = object()
 INPUT_RESUMING = object()
-EMPTY_SEQ = ()
 SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED)
 
-StreamChunk = tuple[tuple[str, ...], str, Any]
-
 
 class StreamProtocol:
     __slots__ = ("modes", "__call__")
@@ -195,6 +197,7 @@ class PregelLoop:
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]],
         stream_keys: Union[str, Sequence[str]],
+        check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
         self.stream = stream
@@ -220,6 +223,11 @@ class PregelLoop:
             self.config = patch_configurable(
                 self.config, {"checkpoint_ns": "", "checkpoint_id": None}
             )
+        if check_subgraphs and self.is_nested and self.checkpointer is not None:
+            if self.config["configurable"]["checkpoint_ns"] in _SEEN_CHECKPOINT_NS:
+                raise MultipleSubgraphsError
+            else:
+                _SEEN_CHECKPOINT_NS.add(self.config["configurable"]["checkpoint_ns"])
         if (
             CONFIG_KEY_CHECKPOINT_MAP in self.config["configurable"]
             and self.config["configurable"].get("checkpoint_ns")
@@ -634,6 +642,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
         super().__init__(
@@ -646,6 +655,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             specs=specs,
             output_keys=output_keys,
             stream_keys=stream_keys,
+            check_subgraphs=check_subgraphs,
             debug=debug,
         )
         self.stack = ExitStack()
@@ -755,6 +765,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
         super().__init__(
@@ -767,6 +778,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             specs=specs,
             output_keys=output_keys,
             stream_keys=stream_keys,
+            check_subgraphs=check_subgraphs,
             debug=debug,
         )
         self.store = AsyncBatchedStore(self.store) if self.store else None

libs/langgraph/langgraph/pregel/messages.py

@@ -24,6 +24,9 @@ Meta = tuple[tuple[str, ...], dict[str, Any]]
 
 
 class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
+    """A callback handler that implements stream_mode=messages.
+    Collects messages from (1) chat model stream events and (2) node outputs."""
+
     def __init__(self, stream: Callable[[StreamChunk], None]):
         self.stream = stream
         self.metadata: dict[UUID, Meta] = {}

libs/langgraph/langgraph/pregel/read.py

@@ -31,6 +31,9 @@ READ_TYPE = Callable[[Union[str, Sequence[str]], bool], Union[Any, dict[str, Any
 
 
 class ChannelRead(RunnableCallable):
+    """Implements the logic for reading state from CONFIG_KEY_READ.
+    Usable both as a runnable as well as a static method to call imperatively."""
+
     channel: Union[str, list[str]]
 
     fresh: bool = False
@@ -108,21 +111,39 @@ DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
 
 
 class PregelNode(Runnable):
+    """A node in a Pregel graph. This won't be invoked as a runnable by the graph
+    itself, but instead acts as a container for the components necessary to make
+    a PregelExecutableTask for a node."""
+
     channels: Union[list[str], Mapping[str, str]]
+    """The channels that will be passed as input to `bound`.
+    If a list, the node will be invoked with the first of that isn't empty.
+    If a dict, the keys are the names of the channels, and the values are the keys
+    to use in the input to `bound`."""
 
     triggers: list[str]
+    """If any of these channels is written to, this node will be triggered in
+    the next step."""
 
     mapper: Optional[Callable[[Any], Any]]
+    """A function to transform the input before passing it to `bound`."""
 
     writers: list[Runnable]
+    """A list of writers that will be executed after `bound`, responsible for
+    taking the output of `bound` and writing it to the appropriate channels."""
 
     bound: Runnable[Any, Any]
+    """The main logic of the node. This will be invoked with the input from 
+    `channels`."""
 
     retry_policy: Optional[RetryPolicy]
+    """The retry policy to use when invoking the node."""
 
     tags: Optional[Sequence[str]]
+    """Tags to attach to the node for tracing."""
 
     metadata: Optional[Mapping[str, Any]]
+    """Metadata to attach to the node for tracing."""
 
     def __init__(
         self,
@@ -151,7 +172,7 @@ class PregelNode(Runnable):
 
     @cached_property
     def flat_writers(self) -> list[Runnable]:
-        """Get writers with optimizations applied."""
+        """Get writers with optimizations applied. Dedupes consecutive ChannelWrites."""
         writers = self.writers.copy()
         while (
             len(writers) > 1
@@ -170,6 +191,7 @@ class PregelNode(Runnable):
 
     @cached_property
     def node(self) -> Optional[Runnable[Any, Any]]:
+        """Get a runnable that combines `bound` and `writers`."""
         writers = self.flat_writers
         if self.bound is DEFAULT_BOUND and not writers:
             return None

libs/langgraph/langgraph/pregel/retry.py

@@ -5,8 +5,8 @@ import time
 from typing import Optional, Sequence
 
 from langgraph.constants import CONFIG_KEY_RESUMING
-from langgraph.errors import GraphInterrupt
-from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
+from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt
+from langgraph.types import PregelExecutableTask, RetryPolicy
 from langgraph.utils.config import patch_configurable
 
 logger = logging.getLogger(__name__)
@@ -71,6 +71,13 @@ def run_with_retry(
             )
             # signal subgraphs to resume (if available)
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
+            # clear checkpoint_ns seen (for subgraph detection)
+            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
+        finally:
+            # clear checkpoint_ns seen (for subgraph detection)
+            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
 
 
 async def arun_with_retry(
@@ -137,3 +144,10 @@ async def arun_with_retry(
             )
             # signal subgraphs to resume (if available)
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
+            # clear checkpoint_ns seen (for subgraph detection)
+            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
+        finally:
+            # clear checkpoint_ns seen (for subgraph detection)
+            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)

libs/langgraph/langgraph/pregel/runner.py

@@ -18,10 +18,14 @@ from langgraph.constants import ERROR, INTERRUPT, NO_WRITES
 from langgraph.errors import GraphDelegate, GraphInterrupt
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
-from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
+from langgraph.types import PregelExecutableTask, RetryPolicy
 
 
 class PregelRunner:
+    """Responsible for executing a set of Pregel tasks concurrently, committing
+    their writes, yielding control to caller when there is output to emit, and
+    interrupting other tasks if appropriate."""
+
     def __init__(
         self,
         *,
@@ -215,6 +219,8 @@ class PregelRunner:
 def _should_stop_others(
     done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Future[Any]]],
 ) -> bool:
+    """Check if any task failed, if so, cancel all other tasks.
+    GraphInterrupts are not considered failures."""
     for fut in done:
         if fut.cancelled():
             return True
@@ -227,6 +233,7 @@ def _should_stop_others(
 def _exception(
     fut: Union[concurrent.futures.Future[Any], asyncio.Future[Any]],
 ) -> Optional[BaseException]:
+    """Return the exception from a future, without raising CancelledError."""
     if fut.cancelled():
         if isinstance(fut, asyncio.Future):
             return asyncio.CancelledError()
@@ -245,6 +252,7 @@ def _panic_or_proceed(
     timeout_exc_cls: Type[Exception] = TimeoutError,
     panic: bool = True,
 ) -> None:
+    """Cancel remaining tasks if any failed, re-raise exception if panic is True."""
     done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     for fut, val in futs.items():

libs/langgraph/langgraph/pregel/types.py

@@ -1,124 +1,25 @@
-from collections import deque
-from typing import Any, Callable, Literal, NamedTuple, Optional, Sequence, Type, Union
+"""Re-export types moved to langgraph.types"""
 
-from langchain_core.runnables import Runnable, RunnableConfig
+from langgraph.types import (
+    All,
+    CachePolicy,
+    PregelExecutableTask,
+    PregelTask,
+    RetryPolicy,
+    StateSnapshot,
+    StreamMode,
+    StreamWriter,
+    default_retry_on,
+)
 
-from langgraph.checkpoint.base import CheckpointMetadata
-from langgraph.constants import Interrupt
-
-
-def default_retry_on(exc: Exception) -> bool:
-    import httpx
-    import requests
-
-    if isinstance(exc, ConnectionError):
-        return True
-    if isinstance(
-        exc,
-        (
-            ValueError,
-            TypeError,
-            ArithmeticError,
-            ImportError,
-            LookupError,
-            NameError,
-            SyntaxError,
-            RuntimeError,
-            ReferenceError,
-            StopIteration,
-            StopAsyncIteration,
-            OSError,
-        ),
-    ):
-        return False
-    if isinstance(exc, httpx.HTTPStatusError):
-        return 500 <= exc.response.status_code < 600
-    if isinstance(exc, requests.HTTPError):
-        return 500 <= exc.response.status_code < 600 if exc.response else True
-    return True
-
-
-class RetryPolicy(NamedTuple):
-    """Configuration for retrying nodes."""
-
-    initial_interval: float = 0.5
-    """Amount of time that must elapse before the first retry occurs. In seconds."""
-    backoff_factor: float = 2.0
-    """Multiplier by which the interval increases after each retry."""
-    max_interval: float = 128.0
-    """Maximum amount of time that may elapse between retries. In seconds."""
-    max_attempts: int = 3
-    """Maximum number of attempts to make before giving up, including the first."""
-    jitter: bool = True
-    """Whether to add random jitter to the interval between retries."""
-    retry_on: Union[
-        Type[Exception], Sequence[Type[Exception]], Callable[[Exception], bool]
-    ] = default_retry_on
-    """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
-
-
-class CachePolicy(NamedTuple):
-    """Configuration for caching nodes."""
-
-    pass
-
-
-class PregelTask(NamedTuple):
-    id: str
-    name: str
-    path: tuple[Union[str, int], ...]
-    error: Optional[Exception] = None
-    interrupts: tuple[Interrupt, ...] = ()
-    state: Union[None, RunnableConfig, "StateSnapshot"] = None
-
-
-class PregelExecutableTask(NamedTuple):
-    name: str
-    input: Any
-    proc: Runnable
-    writes: deque[tuple[str, Any]]
-    config: RunnableConfig
-    triggers: list[str]
-    retry_policy: Optional[RetryPolicy]
-    cache_policy: Optional[CachePolicy]
-    id: str
-    path: tuple[Union[str, int], ...]
-    scheduled: bool = False
-
-
-class StateSnapshot(NamedTuple):
-    """Snapshot of the state of the graph at the beginning of a step."""
-
-    values: Union[dict[str, Any], Any]
-    """Current values of channels"""
-    next: tuple[str, ...]
-    """The name of the node to execute in each task for this step."""
-    config: RunnableConfig
-    """Config used to fetch this snapshot"""
-    metadata: Optional[CheckpointMetadata]
-    """Metadata associated with this snapshot"""
-    created_at: Optional[str]
-    """Timestamp of snapshot creation"""
-    parent_config: Optional[RunnableConfig]
-    """Config used to fetch the parent snapshot, if any"""
-    tasks: tuple[PregelTask, ...]
-    """Tasks to execute in this step. If already attempted, may contain an error."""
-
-
-All = Literal["*"]
-
-StreamMode = Literal["values", "updates", "debug", "messages", "custom"]
-"""How the stream method should emit outputs.
-
-- 'values': Emit all values of the state for each step.
-- 'updates': Emit only the node name(s) and updates
-    that were returned by the node(s) **after** each step.
-- 'debug': Emit debug events for each step.
-- 'messages': Emit LLM messages token-by-token.
-- 'custom': Emit custom output `write: StreamWriter` kwarg of each node.
-"""
-
-StreamWriter = Callable[[Any], None]
-"""Callable that accepts a single argument and writes it to the output stream.
-Always injected into nodes if requested,
-but it's a no-op when not using stream_mode="custom"."""
+__all__ = [
+    "All",
+    "CachePolicy",
+    "PregelExecutableTask",
+    "PregelTask",
+    "RetryPolicy",
+    "StateSnapshot",
+    "StreamMode",
+    "StreamWriter",
+    "default_retry_on",
+]

libs/langgraph/langgraph/pregel/utils.py

@@ -4,7 +4,7 @@ from langgraph.checkpoint.base import ChannelVersions
 def get_new_channel_versions(
     previous_versions: ChannelVersions, current_versions: ChannelVersions
 ) -> ChannelVersions:
-    """Get new channel versions."""
+    """Get subset of current_versions that are newer than previous_versions."""
     if previous_versions:
         version_type = type(next(iter(current_versions.values()), None))
         null_version = version_type()  # type: ignore[misc]

libs/langgraph/langgraph/pregel/validate.py

@@ -3,7 +3,7 @@ from typing import Any, Mapping, Optional, Sequence, Union
 from langgraph.channels.base import BaseChannel
 from langgraph.constants import RESERVED
 from langgraph.pregel.read import PregelNode
-from langgraph.pregel.types import All
+from langgraph.types import All
 
 
 def validate_graph(
@@ -17,7 +17,7 @@ def validate_graph(
 ) -> None:
     for chan in channels:
         if chan in RESERVED:
-            raise ValueError(f"Channel names {RESERVED} are reserved")
+            raise ValueError(f"Channel names {chan} are reserved")
 
     subscribed_channels = set[str]()
     for name, node in nodes.items():

libs/langgraph/langgraph/pregel/write.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 from typing import (
     Any,
     Callable,
@@ -22,30 +21,29 @@ from langgraph.utils.runnable import RunnableCallable
 TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 R = TypeVar("R", bound=Runnable)
 
-
 SKIP_WRITE = object()
 PASSTHROUGH = object()
 
 
 class ChannelWriteEntry(NamedTuple):
     channel: str
+    """Channel name to write to."""
     value: Any = PASSTHROUGH
+    """Value to write, or PASSTHROUGH to use the input."""
     skip_none: bool = False
+    """Whether to skip writing if the value is None."""
     mapper: Optional[Callable] = None
+    """Function to transform the value before writing."""
 
 
 class ChannelWrite(RunnableCallable):
+    """Implements th logic for sending writes to CONFIG_KEY_SEND.
+    Can be used as a runnable or as a static method to call imperatively."""
+
     writes: list[Union[ChannelWriteEntry, Send]]
-    """
-    Sequence of write entries, each of which is a tuple of:
-    - channel name
-    - runnable to map input, or None to use the input, or any other value to use instead
-    - whether to skip writing if the mapped value is None
-    """
+    """Sequence of write entries or Send objects to write."""
     require_at_least_one_of: Optional[Sequence[str]]
-    """
-    If defined, at least one of these channels must be written to.
-    """
+    """If defined, at least one of these channels must be written to."""
 
     def __init__(
         self,
@@ -145,6 +143,7 @@ class ChannelWrite(RunnableCallable):
 
     @staticmethod
     def is_writer(runnable: Runnable) -> bool:
+        """Used by PregelNode to distinguish between writers and other runnables."""
         return (
             isinstance(runnable, ChannelWrite)
             or getattr(runnable, "_is_channel_writer", False) is True
@@ -152,13 +151,9 @@ class ChannelWrite(RunnableCallable):
 
     @staticmethod
     def register_writer(runnable: R) -> R:
+        """Used to mark a runnable as a writer, so that it can be detected by is_writer.
+        Instances of ChannelWrite are automatically marked as writers."""
         # using object.__setattr__ to work around objects that override __setattr__
         # eg. pydantic models and dataclasses
         object.__setattr__(runnable, "_is_channel_writer", True)
         return runnable
-
-
-def _mk_future(val: Any) -> asyncio.Future:
-    fut: asyncio.Future[Any] = asyncio.Future()
-    fut.set_result(val)
-    return fut

libs/langgraph/langgraph/types.py

@@ -0,0 +1,214 @@
+from collections import deque
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Callable,
+    Literal,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+)
+
+from langchain_core.runnables import Runnable, RunnableConfig
+
+from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
+
+All = Literal["*"]
+"""Special value to indicate that graph should interrupt on all nodes."""
+
+Checkpointer = Union[None, Literal[False], BaseCheckpointSaver]
+"""Type of the checkpointer to use for a subgraph. False disables checkpointing,
+even if the parent graph has a checkpointer. None inherits checkpointer."""
+
+StreamMode = Literal["values", "updates", "debug", "messages", "custom"]
+"""How the stream method should emit outputs.
+
+- 'values': Emit all values of the state for each step.
+- 'updates': Emit only the node name(s) and updates
+    that were returned by the node(s) **after** each step.
+- 'debug': Emit debug events for each step.
+- 'messages': Emit LLM messages token-by-token.
+- 'custom': Emit custom output `write: StreamWriter` kwarg of each node.
+"""
+
+StreamWriter = Callable[[Any], None]
+"""Callable that accepts a single argument and writes it to the output stream.
+Always injected into nodes if requested as a keyword argument, but it's a no-op
+when not using stream_mode="custom"."""
+
+
+def default_retry_on(exc: Exception) -> bool:
+    import httpx
+    import requests
+
+    if isinstance(exc, ConnectionError):
+        return True
+    if isinstance(
+        exc,
+        (
+            ValueError,
+            TypeError,
+            ArithmeticError,
+            ImportError,
+            LookupError,
+            NameError,
+            SyntaxError,
+            RuntimeError,
+            ReferenceError,
+            StopIteration,
+            StopAsyncIteration,
+            OSError,
+        ),
+    ):
+        return False
+    if isinstance(exc, httpx.HTTPStatusError):
+        return 500 <= exc.response.status_code < 600
+    if isinstance(exc, requests.HTTPError):
+        return 500 <= exc.response.status_code < 600 if exc.response else True
+    return True
+
+
+class RetryPolicy(NamedTuple):
+    """Configuration for retrying nodes."""
+
+    initial_interval: float = 0.5
+    """Amount of time that must elapse before the first retry occurs. In seconds."""
+    backoff_factor: float = 2.0
+    """Multiplier by which the interval increases after each retry."""
+    max_interval: float = 128.0
+    """Maximum amount of time that may elapse between retries. In seconds."""
+    max_attempts: int = 3
+    """Maximum number of attempts to make before giving up, including the first."""
+    jitter: bool = True
+    """Whether to add random jitter to the interval between retries."""
+    retry_on: Union[
+        Type[Exception], Sequence[Type[Exception]], Callable[[Exception], bool]
+    ] = default_retry_on
+    """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
+
+
+class CachePolicy(NamedTuple):
+    """Configuration for caching nodes."""
+
+    pass
+
+
+@dataclass
+class Interrupt:
+    value: Any
+    when: Literal["during"] = "during"
+
+
+class PregelTask(NamedTuple):
+    id: str
+    name: str
+    path: tuple[Union[str, int], ...]
+    error: Optional[Exception] = None
+    interrupts: tuple[Interrupt, ...] = ()
+    state: Union[None, RunnableConfig, "StateSnapshot"] = None
+
+
+class PregelExecutableTask(NamedTuple):
+    name: str
+    input: Any
+    proc: Runnable
+    writes: deque[tuple[str, Any]]
+    config: RunnableConfig
+    triggers: list[str]
+    retry_policy: Optional[RetryPolicy]
+    cache_policy: Optional[CachePolicy]
+    id: str
+    path: tuple[Union[str, int], ...]
+    scheduled: bool = False
+
+
+class StateSnapshot(NamedTuple):
+    """Snapshot of the state of the graph at the beginning of a step."""
+
+    values: Union[dict[str, Any], Any]
+    """Current values of channels"""
+    next: tuple[str, ...]
+    """The name of the node to execute in each task for this step."""
+    config: RunnableConfig
+    """Config used to fetch this snapshot"""
+    metadata: Optional[CheckpointMetadata]
+    """Metadata associated with this snapshot"""
+    created_at: Optional[str]
+    """Timestamp of snapshot creation"""
+    parent_config: Optional[RunnableConfig]
+    """Config used to fetch the parent snapshot, if any"""
+    tasks: tuple[PregelTask, ...]
+    """Tasks to execute in this step. If already attempted, may contain an error."""
+
+
+class Send:
+    """A message or packet to send to a specific node in the graph.
+
+    The `Send` class is used within a `StateGraph`'s conditional edges to
+    dynamically invoke a node with a custom state at the next step.
+
+    Importantly, the sent state can differ from the core graph's state,
+    allowing for flexible and dynamic workflow management.
+
+    One such example is a "map-reduce" workflow where your graph invokes
+    the same node multiple times in parallel with different states,
+    before aggregating the results back into the main graph's state.
+
+    Attributes:
+        node (str): The name of the target node to send the message to.
+        arg (Any): The state or message to send to the target node.
+
+    Examples:
+        >>> from typing import Annotated
+        >>> import operator
+        >>> class OverallState(TypedDict):
+        ...     subjects: list[str]
+        ...     jokes: Annotated[list[str], operator.add]
+        ...
+        >>> from langgraph.types import Send
+        >>> from langgraph.graph import END, START
+        >>> def continue_to_jokes(state: OverallState):
+        ...     return [Send("generate_joke", {"subject": s}) for s in state['subjects']]
+        ...
+        >>> from langgraph.graph import StateGraph
+        >>> builder = StateGraph(OverallState)
+        >>> builder.add_node("generate_joke", lambda state: {"jokes": [f"Joke about {state['subject']}"]})
+        >>> builder.add_conditional_edges(START, continue_to_jokes)
+        >>> builder.add_edge("generate_joke", END)
+        >>> graph = builder.compile()
+        >>>
+        >>> # Invoking with two subjects results in a generated joke for each
+        >>> graph.invoke({"subjects": ["cats", "dogs"]})
+        {'subjects': ['cats', 'dogs'], 'jokes': ['Joke about cats', 'Joke about dogs']}
+    """
+
+    __slots__ = ("node", "arg")
+
+    node: str
+    arg: Any
+
+    def __init__(self, /, node: str, arg: Any) -> None:
+        """
+        Initialize a new instance of the Send class.
+
+        Args:
+            node (str): The name of the target node to send the message to.
+            arg (Any): The state or message to send to the target node.
+        """
+        self.node = node
+        self.arg = arg
+
+    def __hash__(self) -> int:
+        return hash((self.node, self.arg))
+
+    def __repr__(self) -> str:
+        return f"Send(node={self.node!r}, arg={self.arg!r})"
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, Send)
+            and self.node == value.node
+            and self.arg == value.arg
+        )

libs/langgraph/langgraph/utils/runnable.py

@@ -35,7 +35,7 @@ from langchain_core.tracers._streaming import _StreamingCallbackHandler
 from typing_extensions import TypeGuard
 
 from langgraph.constants import CONFIG_KEY_STREAM_WRITER
-from langgraph.pregel.types import StreamWriter
+from langgraph.types import StreamWriter
 from langgraph.utils.config import (
     ensure_config,
     get_async_callback_manager_for_config,

libs/langgraph/langgraph/version.py

@@ -1,4 +1,4 @@
-"""Main entrypoint into package."""
+"""Exports package version."""
 
 from importlib import metadata
 

libs/langgraph/tests/any_str.py

@@ -1,6 +1,28 @@
 import re
 from typing import Any, Sequence, Union
 
+from typing_extensions import Self
+
+
+class FloatBetween(float):
+    def __new__(cls, min_value: float, max_value: float) -> Self:
+        return super().__new__(cls, min_value)
+
+    def __init__(self, min_value: float, max_value: float) -> None:
+        super().__init__()
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, float)
+            and other >= self.min_value
+            and other <= self.max_value
+        )
+
+    def __hash__(self) -> int:
+        return hash((float(self), self.min_value, self.max_value))
+
 
 class AnyStr(str):
     def __init__(self, prefix: Union[str, re.Pattern] = "") -> None:

libs/langgraph/tests/test_pregel.py

@@ -52,8 +52,8 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.constants import ERROR, PULL, PUSH, Interrupt, Send
-from langgraph.errors import InvalidUpdateError, NodeInterrupt
+from langgraph.constants import ERROR, PULL, PUSH
+from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
 from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
@@ -70,9 +70,9 @@ from langgraph.pregel import (
     StateSnapshot,
 )
 from langgraph.pregel.retry import RetryPolicy
-from langgraph.pregel.types import PregelTask, StreamWriter
 from langgraph.store.memory import MemoryStore
-from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence
+from langgraph.types import Interrupt, PregelTask, Send, StreamWriter
+from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence
 from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS
 from tests.fake_chat import FakeChatModel
 from tests.fake_tracer import FakeTracer
@@ -1861,7 +1861,12 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
         assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100
 
 
-def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None:
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_invoke_join_then_call_other_pregel(
+    mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
 
@@ -1912,6 +1917,17 @@ def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None:
     with ThreadPoolExecutor() as executor:
         assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10
 
+    # add checkpointer
+    app.checkpointer = checkpointer
+    # subgraph is called twice in the same node, through .map(), so raises
+    with pytest.raises(MultipleSubgraphsError):
+        app.invoke([2, 3], {"configurable": {"thread_id": "1"}})
+
+    # set inner graph checkpointer NeverCheckpoint
+    inner_app.checkpointer = False
+    # subgraph still called twice, but checkpointing for inner graph is disabled
+    assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
+
 
 def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -8580,22 +8596,22 @@ def test_stream_subgraphs_during_execution(
     assert chunks == [
         # arrives before "inner" finishes
         (
-            0.0,
+            FloatBetween(0.0, 0.1),
             (
                 (AnyStr("inner:"),),
                 {"inner_1": {"my_key": "got here", "my_other_key": ""}},
             ),
         ),
-        (0.2, ((), {"outer_1": {"my_key": " and parallel"}})),
+        (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})),
         (
-            0.5,
+            FloatBetween(0.5, 0.6),
             (
                 (AnyStr("inner:"),),
                 {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
             ),
         ),
-        (0.5, ((), {"inner": {"my_key": "got here and there"}})),
-        (0.5, ((), {"outer_2": {"my_key": " and back again"}})),
+        (FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})),
+        (FloatBetween(0.5, 0.6), ((), {"outer_2": {"my_key": " and back again"}})),
     ]
 
 

libs/langgraph/tests/test_pregel_async.py

@@ -51,8 +51,8 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.constants import ERROR, PULL, PUSH, Interrupt, Send
-from langgraph.errors import InvalidUpdateError, NodeInterrupt
+from langgraph.constants import ERROR, PULL, PUSH
+from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
@@ -68,9 +68,9 @@ from langgraph.pregel import (
     StateSnapshot,
 )
 from langgraph.pregel.retry import RetryPolicy
-from langgraph.pregel.types import PregelTask, StreamWriter
 from langgraph.store.memory import MemoryStore
-from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence
+from langgraph.types import Interrupt, PregelTask, Send, StreamWriter
+from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence
 from tests.conftest import (
     ALL_CHECKPOINTERS_ASYNC,
     ALL_CHECKPOINTERS_ASYNC_PLUS_NONE,
@@ -2080,7 +2080,10 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
     ]
 
 
-async def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None:
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_invoke_join_then_call_other_pregel(
+    mocker: MockerFixture, checkpointer_name: str
+) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
 
@@ -2133,6 +2136,18 @@ async def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None
         27 for _ in range(10)
     ]
 
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        # add checkpointer
+        app.checkpointer = checkpointer
+        # subgraph is called twice in the same node, through .map(), so raises
+        with pytest.raises(MultipleSubgraphsError):
+            await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}})
+
+        # set inner graph checkpointer NeverCheckpoint
+        inner_app.checkpointer = False
+        # subgraph still called twice, but checkpointing for inner graph is disabled
+        assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
+
 
 async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -7187,22 +7202,22 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None
         assert chunks == [
             # arrives before "inner" finishes
             (
-                0.0,
+                FloatBetween(0.0, 0.1),
                 (
                     (AnyStr("inner:"),),
                     {"inner_1": {"my_key": "got here", "my_other_key": ""}},
                 ),
             ),
-            (0.2, ((), {"outer_1": {"my_key": " and parallel"}})),
+            (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})),
             (
-                0.5,
+                FloatBetween(0.5, 0.6),
                 (
                     (AnyStr("inner:"),),
                     {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
                 ),
             ),
-            (0.5, ((), {"inner": {"my_key": "got here and there"}})),
-            (0.5, ((), {"outer_2": {"my_key": " and back again"}})),
+            (FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})),
+            (FloatBetween(0.5, 0.6), ((), {"outer_2": {"my_key": " and back again"}})),
         ]
 
 

libs/langgraph/tests/test_tracing_interops.py

@@ -5,11 +5,14 @@ from typing import Any, Callable, Tuple, TypedDict, TypeVar
 from unittest.mock import MagicMock
 
 import langsmith as ls
+import pytest
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tracers import LangChainTracer
 
 from langgraph.graph import StateGraph
 
+pytestmark = pytest.mark.anyio
+
 
 def _get_mock_client(**kwargs: Any) -> ls.Client:
     mock_session = MagicMock()
@@ -52,6 +55,7 @@ def wait_for(
     raise ValueError(f"Callable did not return within {total_time}")
 
 
+@pytest.mark.skip("This test times out in CI")
 async def test_nested_tracing():
     lt_py_311 = sys.version_info < (3, 11)
     mock_client = _get_mock_client()
@@ -76,7 +80,7 @@ async def test_nested_tracing():
     child_builder = StateGraph(State)
     child_builder.add_node(child_node)
     child_builder.add_edge("__start__", "child_node")
-    child_graph = child_builder.compile()
+    child_graph = child_builder.compile().with_config(run_name="child_graph")
 
     parent_builder = StateGraph(State)
     parent_builder.add_node(parent_node)
@@ -101,7 +105,7 @@ async def test_nested_tracing():
     # If the callbacks weren't propagated correctly, we'd
     # end up with broken dotted_orders
     parent_run = next(data for data in posts if data["name"] == "parent_node")
-    child_run = next(data for data in posts if data["name"] == "child_node")
+    child_run = next(data for data in posts if data["name"] == "child_graph")
     traceable_run = next(data for data in posts if data["name"] == "some_traceable")
 
     assert child_run["dotted_order"].startswith(traceable_run["dotted_order"])

libs/scheduler-kafka/README.md

@@ -95,7 +95,7 @@ You can pass any of the following values as `kwargs` to either `KafkaOrchestrato
 
 - batch_max_n (int): Maximum number of messages to include in a single batch. Default: 10.
 - batch_max_ms (int): Maximum time in milliseconds to wait for messages to include in a batch. Default: 1000.
-- retry_policy (langgraph.pregel.types.RetryPolicy): Controls which graph-level errors will be retried when processing messages. A good use for this is to retry database errors thrown by the checkpointer. Defaults to None.
+- retry_policy (langgraph.types.RetryPolicy): Controls which graph-level errors will be retried when processing messages. A good use for this is to retry database errors thrown by the checkpointer. Defaults to None.
 
 ### Connection settings
 

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -25,7 +25,6 @@ from langgraph.pregel.executor import (
 )
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.runner import PregelRunner
-from langgraph.pregel.types import RetryPolicy
 from langgraph.scheduler.kafka.retry import aretry, retry
 from langgraph.scheduler.kafka.types import (
     AsyncConsumer,
@@ -38,6 +37,7 @@ from langgraph.scheduler.kafka.types import (
     Sendable,
     Topics,
 )
+from langgraph.types import RetryPolicy
 from langgraph.utils.config import patch_configurable
 
 

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -24,7 +24,6 @@ from langgraph.errors import CheckpointNotLatest, GraphInterrupt
 from langgraph.pregel import Pregel
 from langgraph.pregel.executor import BackgroundExecutor, Submit
 from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
-from langgraph.pregel.types import RetryPolicy
 from langgraph.scheduler.kafka.retry import aretry, retry
 from langgraph.scheduler.kafka.types import (
     AsyncConsumer,
@@ -37,6 +36,7 @@ from langgraph.scheduler.kafka.types import (
     Producer,
     Topics,
 )
+from langgraph.types import RetryPolicy
 from langgraph.utils.config import patch_configurable
 
 
@@ -158,6 +158,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
             specs=graph.channels,
             output_keys=graph.output_channels,
             stream_keys=graph.stream_channels,
+            check_subgraphs=False,
         ) as loop:
             if loop.tick(
                 input_keys=graph.input_channels,
@@ -347,6 +348,7 @@ class KafkaOrchestrator(AbstractContextManager):
             specs=graph.channels,
             output_keys=graph.output_channels,
             stream_keys=graph.stream_channels,
+            check_subgraphs=False,
         ) as loop:
             if loop.tick(
                 input_keys=graph.input_channels,

libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py

@@ -6,7 +6,7 @@ from typing import Awaitable, Callable, Optional
 
 from typing_extensions import ParamSpec
 
-from langgraph.pregel.types import RetryPolicy
+from langgraph.types import RetryPolicy
 
 logger = logging.getLogger(__name__)
 P = ParamSpec("P")