Lint

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]]],

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(

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]:

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."""
 

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