diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 787254687..8c734127c 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -9,6 +9,7 @@ from typing import ( Callable, Generic, Optional, + Sequence, TypeVar, Union, get_args, @@ -38,7 +39,7 @@ from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode def task( *, name: Optional[str] = None, - retry: Optional[RetryPolicy] = None, + retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -55,7 +56,7 @@ def task( __func_or_none__: Optional[Union[Callable[P, Awaitable[T]], Callable[P, T]]] = None, *, name: Optional[str] = None, - retry: Optional[RetryPolicy] = None, + retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], @@ -119,6 +120,10 @@ def task( await add_one.ainvoke([1, 2, 3]) # Returns [2, 3, 4] ``` """ + if isinstance(retry, RetryPolicy): + retry_policies: Optional[Sequence[RetryPolicy]] = (retry,) + else: + retry_policies = retry def decorator( func: Union[Callable[P, Awaitable[T]], Callable[P, T]], @@ -137,7 +142,7 @@ def task( # handle regular functions / partials / callable classes, etc. func.__name__ = name - call_func = functools.partial(call, func, retry=retry) + call_func = functools.partial(call, func, retry=retry_policies) object.__setattr__(call_func, "_is_pregel_task", True) return functools.update_wrapper(call_func, func) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 093faf5d4..0b5019354 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -109,7 +109,7 @@ class StateNodeSpec(NamedTuple): runnable: Runnable metadata: Optional[dict[str, Any]] input: Type[Any] - retry_policy: Optional[RetryPolicy] + retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ @@ -251,7 +251,7 @@ class StateGraph(Graph): *, metadata: Optional[dict[str, Any]] = None, input: Optional[Type[Any]] = None, - retry: Optional[RetryPolicy] = None, + retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None, ) -> Self: """Adds a new node to the state graph. @@ -276,7 +276,7 @@ class StateGraph(Graph): *, metadata: Optional[dict[str, Any]] = None, input: Optional[Type[Any]] = None, - retry: Optional[RetryPolicy] = None, + retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None, ) -> Self: """Adds a new node to the state graph. @@ -300,7 +300,7 @@ class StateGraph(Graph): *, metadata: Optional[dict[str, Any]] = None, input: Optional[Type[Any]] = None, - retry: Optional[RetryPolicy] = None, + retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None, ) -> Self: """Adds a new node to the state graph. @@ -312,7 +312,8 @@ class StateGraph(Graph): action (Optional[RunnableLike]): The action associated with the node. (default: None) metadata (Optional[dict[str, Any]]): The metadata associated with the node. (default: None) input (Optional[Type[Any]]): The input schema for the node. (default: the graph's input schema) - retry (Optional[RetryPolicy]): The policy for retrying the node. (default: None) + retry (Optional[Union[RetryPolicy, Sequence[RetryPolicy]]]): The policy for retrying the node. (default: None) + If a sequence is provided, the first matching policy will be applied. destinations (Optional[Union[dict[str, str], tuple[str, ...]]]): Destinations that indicate where a node can route to. This is useful for edgeless graphs with nodes that return `Command` objects. If a dict is provided, the keys will be used as the target node names and the values will be used as the labels for the edges. diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c9821248e..c5d335d75 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -498,8 +498,8 @@ class Pregel(PregelProtocol): store: Optional[BaseStore] = None """Memory store to use for SharedValues. Defaults to None.""" - retry_policy: Optional[RetryPolicy] = None - """Retry policy to use when running tasks. Set to None to disable.""" + retry_policy: Optional[Sequence[RetryPolicy]] = None + """Retry policies to use when running tasks. Set to None to disable.""" config_type: Optional[Type[Any]] = None @@ -528,7 +528,7 @@ class Pregel(PregelProtocol): debug: Optional[bool] = None, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, - retry_policy: Optional[RetryPolicy] = None, + retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, config_type: Optional[Type[Any]] = None, input_model: Optional[Type[BaseModel]] = None, config: Optional[RunnableConfig] = None, @@ -548,7 +548,10 @@ class Pregel(PregelProtocol): self.debug = debug if debug is not None else get_debug() self.checkpointer = checkpointer self.store = store - self.retry_policy = retry_policy + if isinstance(retry_policy, RetryPolicy): + self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) + else: + self.retry_policy = retry_policy self.config_type = config_type self.input_model = input_model self.config = config diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index e43890498..3679b87c9 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -115,7 +115,7 @@ class Call: func: Callable input: Any - retry: Optional[RetryPolicy] + retry: Optional[Sequence[RetryPolicy]] callbacks: Callbacks def __init__( @@ -123,7 +123,7 @@ class Call: func: Callable, input: Any, *, - retry: Optional[RetryPolicy], + retry: Optional[Sequence[RetryPolicy]], callbacks: Callbacks, ) -> None: self.func = func diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 61a451335..e1856aa80 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -5,7 +5,7 @@ import functools import inspect import sys import types -from typing import Any, Callable, Generator, Generic, Optional, TypeVar, cast +from typing import Any, Callable, Generator, Generic, Optional, Sequence, TypeVar, cast from langchain_core.runnables import Runnable from typing_extensions import ParamSpec @@ -224,7 +224,7 @@ class SyncAsyncFuture(Generic[T], concurrent.futures.Future[T]): def call( func: Callable[P, T], *args: Any, - retry: Optional[RetryPolicy] = None, + retry: Optional[Sequence[RetryPolicy]] = None, **kwargs: Any, ) -> SyncAsyncFuture[T]: config = get_config() diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 05d0c6b60..87e913afa 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -144,8 +144,8 @@ class PregelNode(Runnable): """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.""" + retry_policy: Optional[Sequence[RetryPolicy]] + """The retry policies to use when invoking the node.""" tags: Optional[Sequence[str]] """Tags to attach to the node for tracing.""" @@ -166,7 +166,7 @@ class PregelNode(Runnable): tags: Optional[list[str]] = None, metadata: Optional[Mapping[str, Any]] = None, bound: Optional[Runnable[Any, Any]] = None, - retry_policy: Optional[RetryPolicy] = None, + retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, subgraphs: Optional[Sequence[PregelProtocol]] = None, ) -> None: self.channels = channels @@ -174,7 +174,10 @@ class PregelNode(Runnable): self.mapper = mapper self.writers = writers or [] self.bound = bound if bound is not None else DEFAULT_BOUND - self.retry_policy = retry_policy + if isinstance(retry_policy, RetryPolicy): + self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) + else: + self.retry_policy = retry_policy self.tags = tags self.metadata = metadata if subgraphs is not None: diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 6d0e43b54..1fc3b16ea 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -22,12 +22,11 @@ SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) def run_with_retry( task: PregelExecutableTask, - retry_policy: Optional[RetryPolicy], + retry_policy: Optional[Sequence[RetryPolicy]], configurable: Optional[dict[str, Any]] = None, ) -> 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 config = task.config if configurable is not None: @@ -63,38 +62,39 @@ def run_with_retry( exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") if retry_policy is None: raise + + # Check which retry policy applies to this exception + matching_policy = None + for policy in retry_policy: + if _should_retry_on(policy, exc): + matching_policy = policy + break + + if not matching_policy: + raise + # increment attempts attempts += 1 - # check if we should retry - if isinstance(retry_policy.retry_on, Sequence): - if not isinstance(exc, tuple(retry_policy.retry_on)): - raise - elif isinstance(retry_policy.retry_on, type) and issubclass( - retry_policy.retry_on, Exception - ): - if not isinstance(exc, retry_policy.retry_on): - raise - elif callable(retry_policy.retry_on): - if not retry_policy.retry_on(exc): # type: ignore[call-arg] - raise - else: - raise TypeError( - "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable" - ) # check if we should give up - if attempts >= retry_policy.max_attempts: + if attempts >= matching_policy.max_attempts: raise # sleep before retrying + interval = matching_policy.initial_interval + # Apply backoff factor based on attempt count interval = min( - retry_policy.max_interval, - interval * retry_policy.backoff_factor, + matching_policy.max_interval, + interval * (matching_policy.backoff_factor ** (attempts - 1)), ) - time.sleep( - interval + random.uniform(0, 1) if retry_policy.jitter else interval + + # Apply jitter if configured + sleep_time = ( + interval + random.uniform(0, 1) if matching_policy.jitter else interval ) + time.sleep(sleep_time) + # log the retry logger.info( - f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", + f"Retrying task {task.name} after {sleep_time:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", exc_info=exc, ) # signal subgraphs to resume (if available) @@ -103,13 +103,12 @@ def run_with_retry( async def arun_with_retry( task: PregelExecutableTask, - retry_policy: Optional[RetryPolicy], + retry_policies: Optional[Sequence[RetryPolicy]], stream: bool = False, configurable: Optional[dict[str, Any]] = None, ) -> 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 + retry_policies = task.retry_policy or retry_policies attempts = 0 config = task.config if configurable is not None: @@ -149,41 +148,58 @@ async def arun_with_retry( except Exception as exc: if SUPPORTS_EXC_NOTES: exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") - if retry_policy is None: + if retry_policies is None: raise + + # Check which retry policy applies to this exception + matching_policy = None + for policy in retry_policies: + if _should_retry_on(policy, exc): + matching_policy = policy + break + + if not matching_policy: + raise + # increment attempts attempts += 1 - # check if we should retry - if isinstance(retry_policy.retry_on, Sequence): - if not isinstance(exc, tuple(retry_policy.retry_on)): - raise - elif isinstance(retry_policy.retry_on, type) and issubclass( - retry_policy.retry_on, Exception - ): - if not isinstance(exc, retry_policy.retry_on): - raise - elif callable(retry_policy.retry_on): - if not retry_policy.retry_on(exc): # type: ignore[call-arg] - raise - else: - raise TypeError( - "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable" - ) # check if we should give up - if attempts >= retry_policy.max_attempts: + if attempts >= matching_policy.max_attempts: raise # sleep before retrying + interval = matching_policy.initial_interval + # Apply backoff factor based on attempt count interval = min( - retry_policy.max_interval, - interval * retry_policy.backoff_factor, + matching_policy.max_interval, + interval * (matching_policy.backoff_factor ** (attempts - 1)), ) - await asyncio.sleep( - interval + random.uniform(0, 1) if retry_policy.jitter else interval + + # Apply jitter if configured + sleep_time = ( + interval + random.uniform(0, 1) if matching_policy.jitter else interval ) + await asyncio.sleep(sleep_time) + # log the retry logger.info( - f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", + f"Retrying task {task.name} after {sleep_time:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", exc_info=exc, ) # signal subgraphs to resume (if available) config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) + + +def _should_retry_on(retry_policy: RetryPolicy, exc: Exception) -> bool: + """Check if the given exception should be retried based on the retry policy.""" + if isinstance(retry_policy.retry_on, Sequence): + return isinstance(exc, tuple(retry_policy.retry_on)) + elif isinstance(retry_policy.retry_on, type) and issubclass( + retry_policy.retry_on, Exception + ): + return isinstance(exc, retry_policy.retry_on) + elif callable(retry_policy.retry_on): + return retry_policy.retry_on(exc) # type: ignore[call-arg] + else: + raise TypeError( + "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable" + ) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index cd195c2cc..fc12fc685 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -140,7 +140,7 @@ class PregelRunner: *, reraise: bool = True, timeout: Optional[float] = None, - retry_policy: Optional[RetryPolicy] = None, + retry_policy: Optional[Sequence[RetryPolicy]] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: tasks = tuple(tasks) @@ -269,7 +269,7 @@ class PregelRunner: *, reraise: bool = True, timeout: Optional[float] = None, - retry_policy: Optional[RetryPolicy] = None, + retry_policy: Optional[Sequence[RetryPolicy]] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: loop = asyncio.get_event_loop() @@ -519,7 +519,7 @@ def _call( func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, *, - retry: Optional[RetryPolicy] = None, + retry: Optional[Sequence[RetryPolicy]] = None, callbacks: Callbacks = None, futures: weakref.ref[FuturesDict], schedule_task: weakref.ref[ @@ -600,7 +600,7 @@ def _acall( func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, *, - retry: Optional[RetryPolicy] = None, + retry: Optional[Sequence[RetryPolicy]] = None, callbacks: Callbacks = None, # injected dependencies futures: weakref.ref[FuturesDict], diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index a24ee1ea5..9ce1c2c66 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -75,6 +75,10 @@ def default_retry_on(exc: Exception) -> bool: if isinstance(exc, ConnectionError): return True + 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 if isinstance( exc, ( @@ -93,10 +97,6 @@ def default_retry_on(exc: Exception) -> bool: ), ): 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 @@ -172,7 +172,7 @@ class PregelExecutableTask: writes: deque[tuple[str, Any]] config: RunnableConfig triggers: Sequence[str] - retry_policy: Optional[RetryPolicy] + retry_policy: Optional[Sequence[RetryPolicy]] cache_policy: Optional[CachePolicy] id: str path: tuple[Union[str, int, tuple], ...] diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py new file mode 100644 index 000000000..6ef10d4c7 --- /dev/null +++ b/libs/langgraph/tests/test_retry.py @@ -0,0 +1,342 @@ +from unittest.mock import Mock, patch + +import pytest +from typing_extensions import TypedDict + +from langgraph.graph import START, StateGraph +from langgraph.pregel.retry import _should_retry_on +from langgraph.types import RetryPolicy + + +def test_should_retry_on_single_exception(): + """Test retry with a single exception type.""" + policy = RetryPolicy(retry_on=ValueError) + + # Should retry on ValueError + assert _should_retry_on(policy, ValueError("test error")) is True + + # Should not retry on other exceptions + assert _should_retry_on(policy, TypeError("test error")) is False + assert _should_retry_on(policy, Exception("test error")) is False + + +def test_should_retry_on_sequence_of_exceptions(): + """Test retry with a sequence of exception types.""" + policy = RetryPolicy(retry_on=(ValueError, KeyError)) + + # Should retry on listed exceptions + assert _should_retry_on(policy, ValueError("test error")) is True + assert _should_retry_on(policy, KeyError("test error")) is True + + # Should not retry on other exceptions + assert _should_retry_on(policy, TypeError("test error")) is False + assert _should_retry_on(policy, Exception("test error")) is False + + +def test_should_retry_on_subclass_of_exception(): + """Test retry on subclass of specified exception.""" + + class CustomError(ValueError): + pass + + policy = RetryPolicy(retry_on=ValueError) + + # Should retry on subclass of specified exception + assert _should_retry_on(policy, CustomError("test error")) is True + + +def test_should_retry_on_callable(): + """Test retry with a callable predicate.""" + + # Only retry on ValueError with message containing 'retry' + def should_retry(exc: Exception) -> bool: + return isinstance(exc, ValueError) and "retry" in str(exc) + + policy = RetryPolicy(retry_on=should_retry) + + # Should retry when predicate returns True + assert _should_retry_on(policy, ValueError("please retry this")) is True + + # Should not retry when predicate returns False + assert _should_retry_on(policy, ValueError("other error")) is False + assert _should_retry_on(policy, TypeError("please retry this")) is False + + +def test_should_retry_on_invalid_type(): + """Test retry with an invalid retry_on type.""" + policy = RetryPolicy(retry_on=123) # type: ignore + + with pytest.raises(TypeError, match="retry_on must be an Exception class"): + _should_retry_on(policy, ValueError("test error")) + + +def test_should_retry_on_empty_sequence(): + """Test retry with an empty sequence.""" + policy = RetryPolicy(retry_on=()) + + # Should not retry when sequence is empty + assert _should_retry_on(policy, ValueError("test error")) is False + + +def test_should_retry_default_retry_on(): + """Test the default retry_on function.""" + import httpx + import requests + + # Create a RetryPolicy with default_retry_on + policy = RetryPolicy() + + # Should retry on ConnectionError + assert _should_retry_on(policy, ConnectionError("connection refused")) is True + + # Should not retry on common programming errors + assert _should_retry_on(policy, ValueError("invalid value")) is False + assert _should_retry_on(policy, TypeError("invalid type")) is False + assert _should_retry_on(policy, ArithmeticError("division by zero")) is False + assert _should_retry_on(policy, ImportError("module not found")) is False + assert _should_retry_on(policy, LookupError("key not found")) is False + assert _should_retry_on(policy, NameError("name not defined")) is False + assert _should_retry_on(policy, SyntaxError("invalid syntax")) is False + assert _should_retry_on(policy, RuntimeError("runtime error")) is False + assert _should_retry_on(policy, ReferenceError("weak reference")) is False + assert _should_retry_on(policy, StopIteration()) is False + assert _should_retry_on(policy, StopAsyncIteration()) is False + assert _should_retry_on(policy, OSError("file not found")) is False + + # Should retry on httpx.HTTPStatusError with 5xx status code + response_5xx = Mock() + response_5xx.status_code = 503 + http_error_5xx = httpx.HTTPStatusError( + "server error", request=Mock(), response=response_5xx + ) + assert _should_retry_on(policy, http_error_5xx) is True + + # Should not retry on httpx.HTTPStatusError with 4xx status code + response_4xx = Mock() + response_4xx.status_code = 404 + http_error_4xx = httpx.HTTPStatusError( + "not found", request=Mock(), response=response_4xx + ) + assert _should_retry_on(policy, http_error_4xx) is False + + # Should retry on requests.HTTPError with 5xx status code + response_req_5xx = Mock() + response_req_5xx.status_code = 502 + req_error_5xx = requests.HTTPError("bad gateway") + req_error_5xx.response = response_req_5xx + assert _should_retry_on(policy, req_error_5xx) is True + + # Should not retry on requests.HTTPError with 4xx status code + response_req_4xx = Mock() + response_req_4xx.status_code = 400 + req_error_4xx = requests.HTTPError("bad request") + req_error_4xx.response = response_req_4xx + assert _should_retry_on(policy, req_error_4xx) is False + + # Should retry on requests.HTTPError with no response + req_error_no_resp = requests.HTTPError("connection error") + req_error_no_resp.response = None + assert _should_retry_on(policy, req_error_no_resp) is True + + # Should retry on other exceptions by default + class CustomException(Exception): + pass + + assert _should_retry_on(policy, CustomException("custom error")) is True + + +def test_graph_with_single_retry_policy(): + """Test a simple graph with a single RetryPolicy for a node.""" + + class State(TypedDict): + foo: str + + attempt_count = 0 + + def failing_node(state: State): + nonlocal attempt_count + attempt_count += 1 + if attempt_count < 3: # Fail the first two attempts + raise ValueError("Intentional failure") + return {"foo": "success"} + + def other_node(state: State): + return {"foo": "other_node"} + + # Create a retry policy with specific parameters + retry_policy = RetryPolicy( + max_attempts=3, + initial_interval=0.01, # Short interval for tests + backoff_factor=2.0, + jitter=False, # Disable jitter for predictable timing + retry_on=ValueError, + ) + + # Create and compile the graph + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, retry=retry_policy) + .add_node("other_node", other_node) + .add_edge(START, "failing_node") + .add_edge("failing_node", "other_node") + .compile() + ) + + with patch("time.sleep") as mock_sleep: + result = graph.invoke({"foo": ""}) + + # Verify retry behavior + assert attempt_count == 3 # The node should have been tried 3 times + assert result["foo"] == "other_node" # Final result should be from other_node + + # Verify the sleep intervals + call_args_list = [args[0][0] for args in mock_sleep.call_args_list] + assert call_args_list == [0.01, 0.02] + + +def test_graph_with_jitter_retry_policy(): + """Test a graph with a RetryPolicy that uses jitter.""" + + class State(TypedDict): + foo: str + + attempt_count = 0 + + def failing_node(state): + nonlocal attempt_count + attempt_count += 1 + if attempt_count < 2: # Fail the first attempt + raise ValueError("Intentional failure") + return {"foo": "success"} + + # Create a retry policy with jitter enabled + retry_policy = RetryPolicy( + max_attempts=3, + initial_interval=0.01, + jitter=True, # Enable jitter for randomized backoff + retry_on=ValueError, + ) + + # Create and compile the graph + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, retry=retry_policy) + .add_edge(START, "failing_node") + .compile() + ) + + # Test graph execution with mocked random and sleep + with patch("random.uniform", return_value=0.05) as mock_random, patch( + "time.sleep" + ) as mock_sleep: + result = graph.invoke({"foo": ""}) + + # Verify retry behavior + assert attempt_count == 2 # The node should have been tried twice + assert result["foo"] == "success" + + # Verify jitter was applied + mock_random.assert_called_with(0, 1) # Jitter should use random.uniform(0, 1) + mock_sleep.assert_called_with(0.01 + 0.05) # Sleep should include jitter + + +def test_graph_with_multiple_retry_policies(): + """Test a graph with multiple retry policies for a node.""" + + class State(TypedDict): + foo: str + error_type: str + + attempt_counts = {"value_error": 0, "key_error": 0} + + def failing_node(state): + error_type = state["error_type"] + + if error_type == "value_error": + attempt_counts["value_error"] += 1 + if attempt_counts["value_error"] < 2: + raise ValueError("Value error") + elif error_type == "key_error": + attempt_counts["key_error"] += 1 + if attempt_counts["key_error"] < 3: + raise KeyError("Key error") + + return {"foo": f"recovered_from_{error_type}"} + + # Create multiple retry policies + value_error_policy = RetryPolicy( + max_attempts=2, + initial_interval=0.01, + jitter=False, + retry_on=ValueError, + ) + + key_error_policy = RetryPolicy( + max_attempts=3, + initial_interval=0.02, + jitter=False, + retry_on=KeyError, + ) + + # Create and compile the graph with a list of retry policies + graph = ( + StateGraph(State) + .add_node( + "failing_node", + failing_node, + retry=(value_error_policy, key_error_policy), + ) + .add_edge(START, "failing_node") + .compile() + ) + + # Test ValueError scenario + with patch("time.sleep"): + result_value_error = graph.invoke({"foo": "", "error_type": "value_error"}) + + assert attempt_counts["value_error"] == 2 + assert result_value_error["foo"] == "recovered_from_value_error" + + # Reset attempt counts + attempt_counts = {"value_error": 0, "key_error": 0} + + # Test KeyError scenario + with patch("time.sleep"): + result_key_error = graph.invoke({"foo": "", "error_type": "key_error"}) + + assert attempt_counts["key_error"] == 3 + assert result_key_error["foo"] == "recovered_from_key_error" + + +def test_graph_with_max_attempts_exceeded(): + """Test a graph where max_attempts is exceeded.""" + + class State(TypedDict): + foo: str + + def always_failing_node(state): + raise ValueError("Always fails") + + # Create a retry policy with limited attempts + retry_policy = RetryPolicy( + max_attempts=2, + initial_interval=0.01, + jitter=False, + retry_on=ValueError, + ) + + # Create and compile the graph + graph = ( + StateGraph(State) + .add_node("always_failing", always_failing_node, retry=retry_policy) + .add_edge(START, "always_failing") + .compile() + ) + + # Test graph execution + with patch("time.sleep") as mock_sleep, pytest.raises( + ValueError, match="Always fails" + ): + graph.invoke({"foo": ""}) + + mock_sleep.assert_called_with(0.01)