diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 120d44205..689ac05ab 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -40,7 +40,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[P]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -58,7 +58,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[P]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], @@ -144,7 +144,9 @@ def task( # handle regular functions / partials / callable classes, etc. func.__name__ = name - call_func = functools.partial(call, func, retry=retry_policies, cache=cache) + call_func = functools.partial( + call, func, retry=retry_policies, cache=cache_policy + ) object.__setattr__(call_func, "_is_pregel_task", True) return functools.update_wrapper(call_func, func) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index b8b87b516..00cb3220b 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -113,12 +113,12 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry", "cache", "callbacks") + __slots__ = ("func", "input", "retry", "cache_policy", "callbacks") func: Callable input: tuple[tuple[Any, ...], dict[str, Any]] retry: Optional[Sequence[RetryPolicy]] - cache: Optional[CachePolicy] + cache_policy: Optional[CachePolicy] callbacks: Callbacks def __init__( @@ -127,13 +127,13 @@ class Call: input: tuple[tuple[Any, ...], dict[str, Any]], *, retry: Optional[Sequence[RetryPolicy]], - cache: Optional[CachePolicy], + cache_policy: Optional[CachePolicy], callbacks: Callbacks, ) -> None: self.func = func self.input = input self.retry = retry - self.cache = cache + self.cache_policy = cache_policy self.callbacks = callbacks @@ -639,13 +639,14 @@ def prepare_single_task( if parent_ns else b"", (identifier(call.func) or "__dynamic__").encode(), - call.cache.key(*call.input[0], **call.input[1]), + call.cache_policy.key(*call.input[0], **call.input[1]), ) ) ), - call.cache.ttl, + call.cache_policy.ttl, + call.cache_policy.refresh, ) - if call.cache + if call.cache_policy else None, task_id, task_path, @@ -771,6 +772,7 @@ def prepare_single_task( ) ), proc.cache_policy.ttl, + proc.cache_policy.refresh, ) if proc.cache_policy else None, @@ -915,6 +917,7 @@ def prepare_single_task( ) ), proc.cache_policy.ttl, + proc.cache_policy.refresh, ) if proc.cache_policy else None, diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 5458d3972..bc41eafb5 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -242,12 +242,16 @@ def call( func: Callable[P, T], *args: Any, retry: Optional[Sequence[RetryPolicy]] = None, - cache: Optional[CachePolicy] = None, + cache_policy: Optional[CachePolicy] = None, **kwargs: Any, ) -> SyncAsyncFuture[T]: config = get_config() impl = config[CONF][CONFIG_KEY_CALL] fut = impl( - func, (args, kwargs), retry=retry, cache=cache, callbacks=config["callbacks"] + func, + (args, kwargs), + retry=retry, + cache=cache_policy, + callbacks=config["callbacks"], ) return fut diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 4a5f5c74d..7ab11dbb2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1055,7 +1055,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): if cached := { t.cache_key.key: t for t in self.tasks.values() - if t.cache_key and not t.writes + if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in self.cache.get(cached.keys()).items(): cached[key].writes.extend(values) @@ -1231,7 +1231,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if cached := { t.cache_key.key: t for t in self.tasks.values() - if t.cache_key and not t.writes + if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in (await self.cache.aget(cached.keys())).items(): cached[key].writes.extend(values) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 805523195..42c83c7ca 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -523,7 +523,7 @@ def _call( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, - cache: Optional[CachePolicy] = None, + cache_policy: Optional[CachePolicy] = None, callbacks: Callbacks = None, futures: weakref.ref[FuturesDict], schedule_task: weakref.ref[ @@ -545,7 +545,7 @@ def _call( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, cache=cache, callbacks=callbacks), + Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks), ): if match_cached_writes: match_cached_writes() @@ -630,7 +630,7 @@ def _acall( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, cache=cache, callbacks=callbacks), + Call(func, input, retry=retry, cache_policy=cache, callbacks=callbacks), ): if fut := next( ( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index bd012b729..8b42c0e10 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -140,6 +140,9 @@ class CachePolicy(NamedTuple, Generic[P]): ttl: Optional[int] = None """Time to live for the cache entry in seconds. If None, the entry never expires.""" + refresh: bool = False + """Whether to force a refresh of the cache entry when it is accessed.""" + @dataclasses.dataclass(**_DC_KWARGS) class Interrupt: @@ -191,6 +194,8 @@ class CacheKey(NamedTuple): """Key for the cache entry.""" ttl: Optional[int] """Time to live for the cache entry in seconds.""" + refresh: bool + """Whether to force a refresh of the cache entry when it is accessed.""" @dataclasses.dataclass(**_T_DC_KWARGS) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 93ad840db..9ede99db1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6607,7 +6607,7 @@ def test_multiple_interrupts_functional_cache( counter = 0 - @task(cache=CachePolicy()) + @task(cache_policy=CachePolicy()) def double(x: int) -> int: """Increment the counter.""" nonlocal counter