From 09ca96471402d8594b3e9ca71e06d15a71d9868d Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 21 Nov 2024 08:55:07 -0800
Subject: [PATCH] Wire up retry policy

---
 libs/langgraph/langgraph/func/__init__.py |  2 +-
 libs/langgraph/langgraph/pregel/algo.py   | 18 ++++++++++++++----
 libs/langgraph/langgraph/pregel/runner.py | 12 ++++++++++--
 libs/langgraph/langgraph/types.py         |  6 ++++--
 4 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
index 69e9b6c5d..4db234b81 100644
--- a/libs/langgraph/langgraph/func/__init__.py
+++ b/libs/langgraph/langgraph/func/__init__.py
@@ -48,7 +48,7 @@ def task(
     Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]],
 ]:
     def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]:
-        return update_wrapper(partial(call, func), func)
+        return update_wrapper(partial(call, func, retry=retry), func)
 
     return _task
 
diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 330d46449..4e17dd39d 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -65,7 +65,13 @@ from langgraph.pregel.log import logger
 from langgraph.pregel.manager import ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.store.base import BaseStore
-from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask
+from langgraph.types import (
+    All,
+    LoopProtocol,
+    PregelExecutableTask,
+    PregelTask,
+    RetryPolicy,
+)
 from langgraph.utils.config import merge_configs, patch_config
 
 GetNextVersion = Callable[[Optional[V], BaseChannel], V]
@@ -100,14 +106,18 @@ class PregelTaskWrites(NamedTuple):
 
 
 class Call:
-    __slots__ = ("func", "input")
+    __slots__ = ("func", "input", "retry")
 
     func: Callable
     input: Any
+    retry: Optional[RetryPolicy]
 
-    def __init__(self, func: Callable, input: Any) -> None:
+    def __init__(
+        self, func: Callable, input: Any, *, retry: Optional[RetryPolicy]
+    ) -> None:
         self.func = func
         self.input = input
+        self.retry = retry
 
 
 def should_interrupt(
@@ -545,7 +555,7 @@ def prepare_single_task(
                     },
                 ),
                 triggers,
-                None,
+                call.retry,
                 None,
                 task_id,
                 task_path[:3],
diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py
index 8decbedec..8e3f52ace 100644
--- a/libs/langgraph/langgraph/pregel/runner.py
+++ b/libs/langgraph/langgraph/pregel/runner.py
@@ -143,8 +143,12 @@ class PregelRunner:
             task: PregelExecutableTask,
             func: Callable[[Any], Union[Awaitable[Any], Any]],
             input: Any,
+            *,
+            retry: Optional[RetryPolicy] = None,
         ) -> concurrent.futures.Future[Any]:
-            (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)])
+            (fut,) = writer(
+                task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
+            )
             assert fut is not None, "writer did not return a future for call"
             return fut
 
@@ -320,8 +324,12 @@ class PregelRunner:
             task: PregelExecutableTask,
             func: Callable[[Any], Union[Awaitable[Any], Any]],
             input: Any,
+            *,
+            retry: Optional[RetryPolicy] = None,
         ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
-            (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)])
+            (fut,) = writer(
+                task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
+            )
             assert fut is not None, "writer did not return a future for call"
             if asyncio.iscoroutinefunction(func):
                 return fut
diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py
index 5ce64b52b..c3bf175a2 100644
--- a/libs/langgraph/langgraph/types.py
+++ b/libs/langgraph/langgraph/types.py
@@ -372,12 +372,14 @@ T = TypeVar("T")
 
 
 def call(
-    func: str | Callable[P, T], *args: P.args, **kwargs: P.kwargs
+    func: str | Callable[P, T],
+    *args: P.args,
+    retry: Optional[RetryPolicy] = None,
 ) -> concurrent.futures.Future[T]:
     from langgraph.constants import CONFIG_KEY_CALL
     from langgraph.utils.config import get_configurable
 
     conf = get_configurable()
     impl = conf[CONFIG_KEY_CALL]
-    fut = impl(func, *args, **kwargs)
+    fut = impl(func, *args, retry=retry)
     return fut