diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 396fa6c30..d0d1c461d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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: - pending_writes.extend(writes) + for task in next_tasks: + 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: - pending_writes.extend(writes) + for task in next_tasks: + 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 - for node, _, _, _, config, _, _ in tasks + task.name + 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, ) ) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 5609ffe21..389c59e81 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/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)}" - for name, val, _, _, _, _, _ in next_tasks + f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}" + for task in next_tasks ) ) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index d33741697..71ad967e5 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -104,9 +104,9 @@ def map_output_updates( ] if isinstance(output_channels, str): if updated := [ - (node, value) - for node, _, _, writes, _, _, _ in output_tasks - for chan, value in writes + (task.name, value) + for task in output_tasks + 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, - {chan: value for chan, value in writes if chan in output_channels}, + task.name, + {chan: value for chan, value in task.writes if chan in output_channels}, ) - for node, _, _, writes, _, _, _ in output_tasks - if any(chan in output_channels for chan, _ in writes) + for task in output_tasks + if any(chan in output_channels for chan, _ in task.writes) ]: grouped = defaultdict(list) for node, value in updated: diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index a1b71c046..6e492e2ac 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/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( - 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, + return self.copy( + update=dict( + channels={ + **self.channels, + **{chan: chan for chan in channels}, + } + ), ) def __or__( @@ -202,36 +203,11 @@ class PregelNode(RunnableBindingBase): ], ) -> PregelNode: if ChannelWrite.is_writer(other): - return PregelNode( - channels=self.channels, - triggers=self.triggers, - mapper=self.mapper, - writers=[*self.writers, other], - bound=self.bound, - kwargs=self.kwargs, - config=self.config, - ) + return self.copy(update=dict(writers=[*self.writers, other])) elif self.bound is DEFAULT_BOUND: - return PregelNode( - channels=self.channels, - triggers=self.triggers, - mapper=self.mapper, - writers=self.writers, - bound=coerce_to_runnable(other), - kwargs=self.kwargs, - config=self.config, - ) + return self.copy(update=dict(bound=coerce_to_runnable(other))) else: - return PregelNode( - 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, - ) + return self.copy(update=dict(bound=self.bound | other)) def pipe( self, diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 0fe73a27a..b8da098c4 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/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 -import requests - -from langgraph.pregel.types import PregelExecutableTask +from langgraph.pregel.types import PregelExecutableTask, RetryPolicy 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: diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 7d23eb009..33d79bd9e 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/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