Wire up retry policy

This commit is contained in:
Nuno Campos
2024-12-04 15:39:16 -08:00
parent 0663d46c47
commit 09ca964714
4 changed files with 29 additions and 9 deletions
+1 -1
View File
@@ -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
+14 -4
View File
@@ -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],
+10 -2
View File
@@ -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
+4 -2
View File
@@ -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