mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
Wire up retry policy
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user