pregel: Support retry policy per node

- to expose from StateGraph in future PR

libs/langgraph/langgraph/pregel/__init__.py

@@ -570,6 +570,7 @@ class Pregel(
                 deque(),
                 None,
                 [INTERRUPT],
+                None,
                 str(uuid5(UUID(checkpoint["id"]), INTERRUPT)),
             )
             # execute task
@@ -661,6 +662,7 @@ class Pregel(
                 deque(),
                 None,
                 [INTERRUPT],
+                None,
                 str(uuid5(UUID(checkpoint["id"]), INTERRUPT)),
             )
             # execute task
@@ -1112,8 +1114,8 @@ class Pregel(
 
                     # combine pending writes from all tasks
                     pending_writes = deque[tuple[str, Any]]()
-                    for _, _, _, writes, _, _, _ in next_tasks:
+                    for task in next_tasks:
-                        pending_writes.extend(writes)
+                        pending_writes.extend(task.writes)
 
                     if debug:
                         print_step_writes(
@@ -1560,8 +1562,8 @@ class Pregel(
 
                     # combine pending writes from all tasks
                     pending_writes = deque[tuple[str, Any]]()
-                    for _, _, _, writes, _, _, _ in next_tasks:
+                    for task in next_tasks:
-                        pending_writes.extend(writes)
+                        pending_writes.extend(task.writes)
 
                     if debug:
                         print_step_writes(
@@ -1782,12 +1784,12 @@ def _should_interrupt(
         )
         # and any triggered node is in interrupt_nodes list
         and any(
-            node
+            task.name
-            for node, _, _, _, config, _, _ in tasks
+            for task in tasks
             if (
-                (not config or TAG_HIDDEN not in config.get("tags"))
+                (not task.config or TAG_HIDDEN not in task.config.get("tags"))
                 if interrupt_nodes == "*"
-                else node in interrupt_nodes
+                else task.name in interrupt_nodes
             )
         )
     )
@@ -1936,7 +1938,8 @@ def _prepare_next_tasks(
             logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends")
             continue
         if for_execution:
-            if node := processes[packet.node].get_node():
+            proc = processes[packet.node]
+            if node := proc.get_node():
                 triggers = [TASKS]
                 metadata = {
                     "langgraph_step": step,
@@ -1975,6 +1978,7 @@ def _prepare_next_tasks(
                             },
                         ),
                         triggers,
+                        proc.retry_policy,
                         task_id,
                     )
                 )
@@ -2061,6 +2065,7 @@ def _prepare_next_tasks(
                                 },
                             ),
                             triggers,
+                            proc.retry_policy,
                             task_id,
                         )
                     )

libs/langgraph/langgraph/pregel/debug.py

@@ -66,7 +66,7 @@ def map_debug_tasks(
     step: int, tasks: list[PregelExecutableTask]
 ) -> Iterator[DebugOutputTask]:
     ts = datetime.now(timezone.utc).isoformat()
-    for name, input, _, _, config, triggers, _ in tasks:
+    for name, input, _, _, config, triggers, _, _ in tasks:
         if config is not None and TAG_HIDDEN in config.get("tags", []):
             continue
 
@@ -91,7 +91,7 @@ def map_debug_task_results(
     stream_channels_list: Sequence[str],
 ) -> Iterator[DebugOutputTaskResult]:
     ts = datetime.now(timezone.utc).isoformat()
-    for name, _, _, writes, config, _, _ in tasks:
+    for name, _, _, writes, config, _, _, _ in tasks:
         if config is not None and TAG_HIDDEN in config.get("tags", []):
             continue
 
@@ -137,8 +137,8 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
             f"Starting step {step} with {n_tasks} task{'s' if n_tasks != 1 else ''}:\n"
         )
         + "\n".join(
-            f"- {get_colored_text(name, 'green')} -> {pformat(val)}"
+            f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}"
-            for name, val, _, _, _, _, _ in next_tasks
+            for task in next_tasks
         )
     )
 

libs/langgraph/langgraph/pregel/io.py

@@ -104,9 +104,9 @@ def map_output_updates(
     ]
     if isinstance(output_channels, str):
         if updated := [
-            (node, value)
+            (task.name, value)
-            for node, _, _, writes, _, _, _ in output_tasks
+            for task in output_tasks
-            for chan, value in writes
+            for chan, value in task.writes
             if chan == output_channels
         ]:
             grouped = defaultdict(list)
@@ -119,11 +119,11 @@ def map_output_updates(
     else:
         if updated := [
             (
-                node,
+                task.name,
-                {chan: value for chan, value in writes if chan in output_channels},
+                {chan: value for chan, value in task.writes if chan in output_channels},
             )
-            for node, _, _, writes, _, _, _ in output_tasks
+            for task in output_tasks
-            if any(chan in output_channels for chan, _ in writes)
+            if any(chan in output_channels for chan, _ in task.writes)
         ]:
             grouped = defaultdict(list)
             for node, value in updated:

libs/langgraph/langgraph/pregel/read.py

@@ -16,6 +16,7 @@ from langchain_core.runnables.utils import ConfigurableFieldSpec
 
 from langgraph.constants import CONFIG_KEY_READ
 from langgraph.managed.base import ManagedValueSpec
+from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.write import ChannelWrite
 from langgraph.utils import RunnableCallable
 
@@ -112,6 +113,8 @@ class PregelNode(RunnableBindingBase):
 
     kwargs: Mapping[str, Any] = Field(default_factory=dict)
 
+    retry_policy: Optional[RetryPolicy] = None
+
     def get_writers(self) -> list[Runnable]:
         """Get writers with optimizations applied."""
         writers = self.writers.copy()
@@ -155,6 +158,7 @@ class PregelNode(RunnableBindingBase):
         bound: Optional[Runnable[Any, Any]] = None,
         kwargs: Optional[Mapping[str, Any]] = None,
         config: Optional[RunnableConfig] = None,
+        retry_policy: Optional[RetryPolicy] = None,
         **other_kwargs: Any,
     ) -> None:
         super().__init__(
@@ -164,6 +168,7 @@ class PregelNode(RunnableBindingBase):
             writers=writers or [],
             bound=bound or DEFAULT_BOUND,
             kwargs=kwargs or {},
+            retry_policy=retry_policy,
             config=merge_configs(
                 config, {"tags": tags or [], "metadata": metadata or {}}
             ),
@@ -180,17 +185,13 @@ class PregelNode(RunnableBindingBase):
         assert isinstance(
             self.channels, dict
         ), "all channels must be named when using .join()"
-        return PregelNode(
+        return self.copy(
-            channels={
+            update=dict(
-                **self.channels,
+                channels={
-                **{chan: chan for chan in channels},
+                    **self.channels,
-            },
+                    **{chan: chan for chan in channels},
-            triggers=self.triggers,
+                }
-            mapper=self.mapper,
+            ),
-            writers=self.writers,
-            bound=self.bound,
-            kwargs=self.kwargs,
-            config=self.config,
         )
 
     def __or__(
@@ -202,36 +203,11 @@ class PregelNode(RunnableBindingBase):
         ],
     ) -> PregelNode:
         if ChannelWrite.is_writer(other):
-            return PregelNode(
+            return self.copy(update=dict(writers=[*self.writers, other]))
-                channels=self.channels,
-                triggers=self.triggers,
-                mapper=self.mapper,
-                writers=[*self.writers, other],
-                bound=self.bound,
-                kwargs=self.kwargs,
-                config=self.config,
-            )
         elif self.bound is DEFAULT_BOUND:
-            return PregelNode(
+            return self.copy(update=dict(bound=coerce_to_runnable(other)))
-                channels=self.channels,
-                triggers=self.triggers,
-                mapper=self.mapper,
-                writers=self.writers,
-                bound=coerce_to_runnable(other),
-                kwargs=self.kwargs,
-                config=self.config,
-            )
         else:
-            return PregelNode(
+            return self.copy(update=dict(bound=self.bound | other))
-                channels=self.channels,
-                triggers=self.triggers,
-                mapper=self.mapper,
-                writers=self.writers,
-                # delegate to __or__ in self.bound
-                bound=self.bound | other,
-                kwargs=self.kwargs,
-                config=self.config,
-            )
 
     def pipe(
         self,

libs/langgraph/langgraph/pregel/retry.py

@@ -2,68 +2,19 @@ import asyncio
 import logging
 import random
 import time
-from typing import Callable, NamedTuple, Optional, Union
+from typing import Optional
 
-import httpx
+from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
-import requests
-
-from langgraph.pregel.types import PregelExecutableTask
 
 logger = logging.getLogger(__name__)
 
 
-def default_retry_on(exc: Exception) -> bool:
-    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 = 10
-    """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[
-        tuple[Exception, ...], Callable[[Exception], bool]
-    ] = default_retry_on
-    """List of exceptions that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
-
-
 def run_with_retry(
     task: PregelExecutableTask,
     retry_policy: Optional[RetryPolicy],
 ) -> None:
     """Run a task with retries."""
+    retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     while True:
@@ -108,6 +59,7 @@ async def arun_with_retry(
     stream: bool = False,
 ) -> None:
     """Run a task asynchronously with retries."""
+    retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     while True:

libs/langgraph/langgraph/pregel/types.py

@@ -1,11 +1,61 @@
 from collections import deque
-from typing import Any, Literal, NamedTuple, Optional, Union
+from typing import Any, Callable, Literal, NamedTuple, Optional, Union
 
 from langchain_core.runnables import Runnable, RunnableConfig
 
 from langgraph.checkpoint.base import CheckpointMetadata
 
 
+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 = 10
+    """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[
+        tuple[Exception, ...], Callable[[Exception], bool]
+    ] = default_retry_on
+    """List of exceptions that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
+
+
 class PregelTaskDescription(NamedTuple):
     name: str
     input: Any
@@ -18,6 +68,7 @@ class PregelExecutableTask(NamedTuple):
     writes: deque[tuple[str, Any]]
     config: RunnableConfig
     triggers: list[str]
+    retry_policy: Optional[RetryPolicy]
     id: str