diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index b2a4081fb..29338b02b 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_policy: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[Callable[P, str | bytes]]] = 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_policy: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[Callable[P, str | bytes]]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index ff94ec0f6..225b0ccd3 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -599,7 +599,7 @@ def prepare_single_task( writes: deque[tuple[str, Any]] = deque() cache_policy = call.cache_policy or cache_policy if cache_policy: - args_key = cache_policy.key(*call.input[0], **call.input[1]) + args_key = cache_policy.key_func(*call.input[0], **call.input[1]) cache_key: Optional[CacheKey] = CacheKey( xxh3_128_hexdigest( b"".join( @@ -729,7 +729,7 @@ def prepare_single_task( writes = deque() cache_policy = proc.cache_policy or cache_policy if cache_policy: - args_key = cache_policy.key(packet.arg) + args_key = cache_policy.key_func(packet.arg) cache_key = CacheKey( xxh3_128_hexdigest( b"".join( @@ -868,7 +868,7 @@ def prepare_single_task( writes = deque() cache_policy = proc.cache_policy or cache_policy if cache_policy: - args_key = cache_policy.key(val) + args_key = cache_policy.key_func(val) cache_key = CacheKey( xxh3_128_hexdigest( b"".join( diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index dcf32f96e..916abef36 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -142,7 +142,9 @@ class PregelRunner: timeout: Optional[float] = None, retry_policy: Optional[Sequence[RetryPolicy]] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, - match_cached_writes: Optional[Callable[[], None]] = None, + match_cached_writes: Optional[ + Callable[[], Sequence[PregelExecutableTask]] + ] = None, ) -> Iterator[None]: tasks = tuple(tasks) futures = FuturesDict( @@ -529,7 +531,7 @@ def _call( [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ] ], - match_cached_writes: Optional[Callable[[], None]], + match_cached_writes: Optional[Callable[[], Sequence[PregelExecutableTask]]], submit: weakref.ref[Submit], reraise: bool, ) -> concurrent.futures.Future[Any]: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 79da90b67..4b7e36ee0 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -11,7 +11,6 @@ from typing import ( Literal, NamedTuple, Optional, - ParamSpec, TypeVar, Union, cast, @@ -124,16 +123,16 @@ class RetryPolicy(NamedTuple): """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" -P = ParamSpec("P") +KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes]) -class CachePolicy(NamedTuple, Generic[P]): +class CachePolicy(NamedTuple, Generic[KeyFuncT]): """Configuration for caching nodes. !!! version-added "Added in version 0.2.24." """ - key: Callable[P, str | bytes] = default_cache_key + key_func: KeyFuncT = default_cache_key # type: ignore[assignment] """Function to generate a cache key from the node's input. Defaults to hashing the input with pickle.""" diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py index ddae896c8..01ee2b2ea 100644 --- a/libs/langgraph/langgraph/utils/cache.py +++ b/libs/langgraph/langgraph/utils/cache.py @@ -18,7 +18,7 @@ def _freeze(obj: Any) -> Hashable: return obj # strings, ints, dataclasses with frozen=True, etc. -def default_cache_key(*args: Any, **kwargs: Any) -> bytes: +def default_cache_key(*args: Any, **kwargs: Any) -> str | bytes: """Default cache key function that uses the arguments and keyword arguments to generate a hashable key.""" import pickle