Cache nodes/tasks (#4486)

- BaseCache interface defines the base class for cache storage adapters
- FileCache implements BaseCache with filesystem-backed storage
- Provide default cache key implementation which hashes args with pickle
- Update PregelExecutableTask with cache_key property for tasks that
opt-in to caching
- Update PregelLoop, PregelRunner to get/set from cache as appropriate

TODO
- [x]  Call match_cached_writes in async PregelRunner
- [ ]  Implement RedisCache to use in LGP
- [x]  Add more tests

libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import asyncio
+import datetime
+import sqlite3
+import threading
+from collections.abc import Mapping, Sequence
+
+from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class SqliteCache(BaseCache[ValueT]):
+    """File-based cache using SQLite."""
+
+    def __init__(
+        self,
+        *,
+        path: str,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        """Initialize the cache with a file path."""
+        super().__init__(serde=serde)
+        # SQLite backing store
+        self._conn = sqlite3.connect(
+            path,
+            check_same_thread=False,
+        )
+        # Serialize access to the shared connection across threads
+        self._lock = threading.RLock()
+        # Better concurrency & atomicity
+        self._conn.execute("PRAGMA journal_mode=WAL;")
+        # Schema: key -> (expiry, encoding, value)
+        self._conn.execute(
+            """CREATE TABLE IF NOT EXISTS cache (
+                ns TEXT,
+                key TEXT,
+                expiry REAL,
+                encoding TEXT NOT NULL,
+                val BLOB NOT NULL,
+                PRIMARY KEY (ns, key)
+            )"""
+        )
+        self._conn.commit()
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        with self._lock, self._conn:
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            if not keys:
+                return {}
+            placeholders = ",".join("(?, ?)" for _ in keys)
+            params: list[str] = []
+            for ns_tuple, key in keys:
+                params.extend((",".join(ns_tuple), key))
+            cursor = self._conn.execute(
+                f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})",
+                tuple(params),
+            )
+            values: dict[FullKey, ValueT] = {}
+            rows = cursor.fetchall()
+            for ns, key, expiry, encoding, raw in rows:
+                if expiry is not None and now > expiry:
+                    # purge expired entry
+                    self._conn.execute(
+                        "DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key)
+                    )
+                    continue
+                values[(tuple(ns.split(",")), key)] = self.serde.loads_typed(
+                    (encoding, raw)
+                )
+            return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return await asyncio.to_thread(self.get, keys)
+
+    def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys and TTLs."""
+        with self._lock, self._conn:
+            now = datetime.datetime.now(datetime.timezone.utc)
+            for key, (value, ttl) in mapping.items():
+                if ttl is not None:
+                    delta = datetime.timedelta(seconds=ttl)
+                    expiry: float | None = (now + delta).timestamp()
+                else:
+                    expiry = None
+                encoding, raw = self.serde.dumps_typed(value)
+                self._conn.execute(
+                    "INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)",
+                    (",".join(key[0]), key[1], expiry, encoding, raw),
+                )
+
+    async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys and TTLs."""
+        await asyncio.to_thread(self.set, mapping)
+
+    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        with self._lock, self._conn:
+            if namespaces is None:
+                self._conn.execute("DELETE FROM cache")
+            else:
+                placeholders = ",".join("?" for _ in namespaces)
+                self._conn.execute(
+                    f"DELETE FROM cache WHERE (ns) IN ({placeholders})",
+                    tuple(",".join(key) for key in namespaces),
+                )
+
+    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Asynchronously delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        await asyncio.to_thread(self.clear, namespaces)
+
+    def __del__(self) -> None:
+        try:
+            self._conn.close()
+        except Exception:
+            pass

libs/checkpoint/langgraph/cache/base/__init__.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import Generic, Sequence, TypeVar
+
+from langgraph.checkpoint.serde.base import SerializerProtocol
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+ValueT = TypeVar("ValueT")
+Namespace = tuple[str, ...]
+FullKey = tuple[Namespace, str]
+
+
+class BaseCache(ABC, Generic[ValueT]):
+    """Base class for a cache."""
+
+    serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=True)
+
+    def __init__(self, *, serde: SerializerProtocol | None = None) -> None:
+        """Initialize the cache with a serializer."""
+        self.serde = serde or self.serde
+
+    @abstractmethod
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+
+    @abstractmethod
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+
+    @abstractmethod
+    def set(self, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys and TTLs."""
+
+    @abstractmethod
+    async def aset(self, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys and TTLs."""
+
+    @abstractmethod
+    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+
+    @abstractmethod
+    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Asynchronously delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""

libs/checkpoint/langgraph/cache/memory/__init__.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import datetime
+import threading
+from collections.abc import Mapping, Sequence
+
+from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class InMemoryCache(BaseCache[ValueT]):
+    def __init__(self, *, serde: SerializerProtocol | None = None):
+        super().__init__(serde=serde)
+        self._cache: dict[Namespace, dict[str, tuple[str, bytes, float | None]]] = {}
+        self._lock = threading.RLock()
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        with self._lock:
+            if not keys:
+                return {}
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            values: dict[FullKey, ValueT] = {}
+            for ns_tuple, key in keys:
+                ns = Namespace(ns_tuple)
+                if ns in self._cache and key in self._cache[ns]:
+                    enc, val, expiry = self._cache[ns][key]
+                    if expiry is None or now < expiry:
+                        values[(ns, key)] = self.serde.loads_typed((enc, val))
+                    else:
+                        del self._cache[ns][key]
+            return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return self.get(keys)
+
+    def set(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys."""
+        with self._lock:
+            now = datetime.datetime.now(datetime.timezone.utc)
+            for (ns, key), (value, ttl) in keys.items():
+                if ttl is not None:
+                    delta = datetime.timedelta(seconds=ttl)
+                    expiry: float | None = (now + delta).timestamp()
+                else:
+                    expiry = None
+                if ns not in self._cache:
+                    self._cache[ns] = {}
+                self._cache[ns][key] = (
+                    *self.serde.dumps_typed(value),
+                    expiry,
+                )
+
+    async def aset(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys."""
+        self.set(keys)
+
+    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        with self._lock:
+            if namespaces is None:
+                self._cache.clear()
+            else:
+                for ns in namespaces:
+                    if ns in self._cache:
+                        del self._cache[ns]
+
+    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:
+        """Asynchronously delete the cached values for the given namespaces.
+        If no namespaces are provided, clear all cached values."""
+        self.clear(namespaces)

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -68,8 +68,9 @@ class InMemorySaver(
             str, dict[str, tuple[tuple[str, bytes], tuple[str, bytes], Optional[str]]]
         ],
     ]
+    # (thread ID, checkpoint NS, checkpoint ID) -> (task ID, write idx)
     writes: defaultdict[
-        tuple[str, str, str],  # thread ID, checkpoint NS, checkpoint ID
+        tuple[str, str, str],
         dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]],
     ]
     blobs: dict[

libs/checkpoint/langgraph/checkpoint/serde/base.py

@@ -1,7 +1,17 @@
+from __future__ import annotations
+
 from typing import Any, Protocol
 
 
-class SerializerProtocol(Protocol):
+class UntypedSerializerProtocol(Protocol):
+    """Protocol for serialization and deserialization of objects."""
+
+    def dumps(self, obj: Any) -> bytes: ...
+
+    def loads(self, data: bytes) -> Any: ...
+
+
+class SerializerProtocol(UntypedSerializerProtocol, Protocol):
     """Protocol for serialization and deserialization of objects.
 
     - `dumps`: Serialize an object to bytes.
@@ -12,17 +22,13 @@ class SerializerProtocol(Protocol):
     Valid implementations include the `pickle`, `json` and `orjson` modules.
     """
 
-    def dumps(self, obj: Any) -> bytes: ...
-
     def dumps_typed(self, obj: Any) -> tuple[str, bytes]: ...
 
-    def loads(self, data: bytes) -> Any: ...
-
     def loads_typed(self, data: tuple[str, bytes]) -> Any: ...
 
 
 class SerializerCompat(SerializerProtocol):
-    def __init__(self, serde: SerializerProtocol) -> None:
+    def __init__(self, serde: UntypedSerializerProtocol) -> None:
         self.serde = serde
 
     def dumps(self, obj: Any) -> bytes:
@@ -38,7 +44,9 @@ class SerializerCompat(SerializerProtocol):
         return self.serde.loads(data[1])
 
 
-def maybe_add_typed_methods(serde: SerializerProtocol) -> SerializerProtocol:
+def maybe_add_typed_methods(
+    serde: SerializerProtocol | UntypedSerializerProtocol,
+) -> SerializerProtocol:
     """Wrap serde old serde implementations in a class with loads_typed and dumps_typed for backwards compatibility."""
 
     if not hasattr(serde, "loads_typed") or not hasattr(serde, "dumps_typed"):

libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

@@ -3,6 +3,7 @@ import decimal
 import importlib
 import json
 import pathlib
+import pickle
 import re
 from collections import deque
 from collections.abc import Sequence
@@ -37,8 +38,12 @@ class JsonPlusSerializer(SerializerProtocol):
     """Serializer that uses ormsgpack, with a fallback to extended JSON serializer."""
 
     def __init__(
-        self, *, __unpack_ext_hook__: Optional[Callable[[int, bytes], Any]] = None
+        self,
+        *,
+        pickle_fallback: bool = False,
+        __unpack_ext_hook__: Optional[Callable[[int, bytes], Any]] = None,
     ) -> None:
+        self.pickle_fallback = pickle_fallback
         self._unpack_ext_hook = (
             __unpack_ext_hook__
             if __unpack_ext_hook__ is not None
@@ -209,6 +214,8 @@ class JsonPlusSerializer(SerializerProtocol):
             except ormsgpack.MsgpackEncodeError as exc:
                 if "valid UTF-8" in str(exc):
                     return "json", self.dumps(obj)
+                elif self.pickle_fallback:
+                    return "pickle", pickle.dumps(obj)
                 raise exc
 
     def loads(self, data: bytes) -> Any:
@@ -228,6 +235,8 @@ class JsonPlusSerializer(SerializerProtocol):
             return ormsgpack.unpackb(
                 data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
             )
+        elif self.pickle_fallback and type_ == "pickle":
+            return pickle.loads(data_)
         else:
             raise NotImplementedError(f"Unknown serialization type: {type_}")
 

libs/langgraph/langgraph/constants.py

@@ -46,6 +46,10 @@ TASKS = sys.intern("__pregel_tasks")
 RETURN = sys.intern("__return__")
 # for writes of a task where we simply record the return value
 
+# --- Reserved cache namespaces ---
+CACHE_NS_WRITES = sys.intern("__pregel_ns_writes")
+# cache namespace for node writes
+
 # --- Reserved config.configurable keys ---
 CONFIG_KEY_SEND = sys.intern("__pregel_send")
 # holds the `write` function that accepts writes to state/edges/reserved keys
@@ -61,6 +65,8 @@ CONFIG_KEY_STREAM_WRITER = sys.intern("__pregel_stream_writer")
 # holds a `StreamWriter` for stream_mode=custom
 CONFIG_KEY_STORE = sys.intern("__pregel_store")
 # holds a `BaseStore` made available to managed values
+CONFIG_KEY_CACHE = sys.intern("__pregel_cache")
+# holds a `BaseCache` made available to subgraphs
 CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming")
 # holds a boolean indicating if subgraphs should resume from a previous checkpoint
 CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id")

libs/langgraph/langgraph/func/__init__.py

@@ -16,10 +16,11 @@ from typing import (
     overload,
 )
 
+from langgraph.cache.base import BaseCache
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.constants import END, PREVIOUS, START
+from langgraph.constants import CACHE_NS_WRITES, END, PREVIOUS, START
 from langgraph.pregel import Pregel
 from langgraph.pregel.call import (
     P,
@@ -27,11 +28,55 @@ from langgraph.pregel.call import (
     T,
     call,
     get_runnable_for_entrypoint,
+    identifier,
 )
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
-from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode
+from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode
+
+
+class TaskFunction(Generic[P, T]):
+    def __init__(
+        self,
+        func: Callable[P, T],
+        *,
+        retry: Optional[Sequence[RetryPolicy]] = (),
+        cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None,
+        name: Optional[str] = None,
+    ) -> None:
+        if name is not None:
+            if hasattr(func, "__func__"):
+                # handle class methods
+                # NOTE: we're modifying the instance method to avoid modifying
+                # the original class method in case it's shared across multiple tasks
+                instance_method = functools.partial(func.__func__, func.__self__)  # type: ignore [attr-defined]
+                instance_method.__name__ = name  # type: ignore [attr-defined]
+                func = instance_method
+            else:
+                # handle regular functions / partials / callable classes, etc.
+                func.__name__ = name
+        self.func = func
+        self.retry = retry
+        self.cache_policy = cache_policy
+        functools.update_wrapper(self, func)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> SyncAsyncFuture[T]:
+        return call(
+            self.func, retry=self.retry, cache_policy=self.cache_policy, *args, **kwargs
+        )
+
+    def clear_cache(self, cache: BaseCache) -> None:
+        """Clear the cache for this task."""
+        if self.cache_policy is not None:
+            cache.clear(((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),))
+
+    async def aclear_cache(self, cache: BaseCache) -> None:
+        """Clear the cache for this task."""
+        if self.cache_policy is not None:
+            await cache.aclear(
+                ((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),)
+            )
 
 
 @overload
@@ -39,16 +84,17 @@ def task(
     *,
     name: Optional[str] = None,
     retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+    cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None,
 ) -> Callable[
     [Union[Callable[P, Awaitable[T]], Callable[P, T]]],
-    Callable[P, SyncAsyncFuture[T]],
+    TaskFunction[P, T],
 ]: ...
 
 
 @overload
 def task(
     __func_or_none__: Union[Callable[P, Awaitable[T]], Callable[P, T]],
-) -> Callable[P, SyncAsyncFuture[T]]: ...
+) -> TaskFunction[P, T]: ...
 
 
 def task(
@@ -56,12 +102,13 @@ def task(
     *,
     name: Optional[str] = None,
     retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+    cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None,
 ) -> Union[
     Callable[
         [Union[Callable[P, Awaitable[T]], Callable[P, T]]],
-        Callable[P, SyncAsyncFuture[T]],
+        TaskFunction[P, T],
     ],
-    Callable[P, SyncAsyncFuture[T]],
+    TaskFunction[P, T],
 ]:
     """Define a LangGraph task using the `task` decorator.
 
@@ -129,21 +176,9 @@ def task(
     ) -> Union[
         Callable[P, concurrent.futures.Future[T]], Callable[P, asyncio.Future[T]]
     ]:
-        if name is not None:
-            if hasattr(func, "__func__"):
-                # handle class methods
-                # NOTE: we're modifying the instance method to avoid modifying
-                # the original class method in case it's shared across multiple tasks
-                instance_method = functools.partial(func.__func__, func.__self__)  # type: ignore [union-attr]
-                instance_method.__name__ = name  # type: ignore [attr-defined]
-                func = instance_method
-            else:
-                # handle regular functions / partials / callable classes, etc.
-                func.__name__ = name
-
-        call_func = functools.partial(call, func, retry=retry_policies)
-        object.__setattr__(call_func, "_is_pregel_task", True)
-        return functools.update_wrapper(call_func, func)
+        return TaskFunction(
+            func, retry=retry_policies, cache_policy=cache_policy, name=name
+        )
 
     if __func_or_none__ is not None:
         return decorator(__func_or_none__)
@@ -316,11 +351,17 @@ class entrypoint:
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
         store: Optional[BaseStore] = None,
+        cache: Optional[BaseCache] = None,
         config_schema: Optional[type[Any]] = None,
+        cache_policy: Optional[CachePolicy] = None,
+        retry: Union[RetryPolicy, Sequence[RetryPolicy]] = (),
     ) -> None:
         """Initialize the entrypoint decorator."""
         self.checkpointer = checkpointer
         self.store = store
+        self.cache = cache
+        self.cache_policy = cache_policy
+        self.retry = retry
         self.config_schema = config_schema
 
     @dataclass(**_DC_KWARGS)
@@ -450,5 +491,8 @@ class entrypoint:
             stream_eager=True,
             checkpointer=self.checkpointer,
             store=self.store,
+            cache=self.cache,
+            cache_policy=self.cache_policy,
+            retry_policy=self.retry,
             config_type=self.config_schema,
         )

libs/langgraph/langgraph/graph/graph.py

@@ -14,6 +14,7 @@ from typing import (
 from langchain_core.runnables import Runnable
 from typing_extensions import Self
 
+from langgraph.cache.base import BaseCache
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.constants import (
     EMPTY_SEQ,
@@ -28,6 +29,7 @@ from langgraph.graph.branch import Branch
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 from langgraph.types import All, Checkpointer
 from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
 
@@ -316,6 +318,9 @@ class Graph:
         interrupt_after: Optional[Union[All, list[str]]] = None,
         debug: bool = False,
         name: Optional[str] = None,
+        *,
+        cache: Optional[BaseCache] = None,
+        store: Optional[BaseStore] = None,
     ) -> "CompiledGraph":
         """Compiles the graph into a `CompiledGraph` object.
 
@@ -364,6 +369,8 @@ class Graph:
             auto_validate=False,
             debug=debug,
             name=name or "LangGraph",
+            cache=cache,
+            store=store,
         )
 
         # attach nodes, edges, and branches

libs/langgraph/langgraph/graph/state.py

@@ -26,6 +26,7 @@ from pydantic import BaseModel
 from typing_extensions import Self
 
 from langgraph._api.deprecation import LangGraphDeprecationWarning
+from langgraph.cache.base import BaseCache
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.dynamic_barrier_value import (
@@ -79,7 +80,7 @@ from langgraph.pregel.write import (
     ChannelWriteTupleEntry,
 )
 from langgraph.store.base import BaseStore
-from langgraph.types import All, Checkpointer, Command, RetryPolicy
+from langgraph.types import All, CachePolicy, Checkpointer, Command, RetryPolicy
 from langgraph.utils.fields import get_field_default, get_update_as_tuples
 from langgraph.utils.pydantic import create_model
 from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
@@ -113,6 +114,7 @@ class StateNodeSpec(NamedTuple):
     metadata: Optional[dict[str, Any]]
     input: type[Any]
     retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]]
+    cache_policy: Optional[CachePolicy]
     ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ
     defer: bool = False
 
@@ -259,6 +261,7 @@ class StateGraph(Graph):
         metadata: Optional[dict[str, Any]] = None,
         input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+        cache_policy: Optional[CachePolicy] = None,
         destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None,
     ) -> Self:
         """Add a new node to the state graph.
@@ -276,6 +279,7 @@ class StateGraph(Graph):
         metadata: Optional[dict[str, Any]] = None,
         input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+        cache_policy: Optional[CachePolicy] = None,
         destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None,
     ) -> Self:
         """Add a new node to the state graph."""
@@ -290,6 +294,7 @@ class StateGraph(Graph):
         metadata: Optional[dict[str, Any]] = None,
         input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+        cache_policy: Optional[CachePolicy] = None,
         destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None,
     ) -> Self:
         """Add a new node to the state graph.
@@ -303,6 +308,7 @@ class StateGraph(Graph):
             input: The input schema for the node. (default: the graph's input schema)
             retry: The policy for retrying the node. (default: None)
                 If a sequence is provided, the first matching policy will be applied.
+            cache_policy: The cache policy for the node. (default: None)
             destinations: Destinations that indicate where a node can route to.
                 This is useful for edgeless graphs with nodes that return `Command` objects.
                 If a dict is provided, the keys will be used as the target node names and the values will be used as the labels for the edges.
@@ -431,6 +437,7 @@ class StateGraph(Graph):
             metadata,
             input=input or self.schema,
             retry_policy=retry,
+            cache_policy=cache_policy,
             ends=ends,
             defer=defer,
         )
@@ -571,6 +578,7 @@ class StateGraph(Graph):
         self,
         checkpointer: Checkpointer = None,
         *,
+        cache: Optional[BaseCache] = None,
         store: Optional[BaseStore] = None,
         interrupt_before: Optional[Union[All, list[str]]] = None,
         interrupt_after: Optional[Union[All, list[str]]] = None,
@@ -655,6 +663,7 @@ class StateGraph(Graph):
             auto_validate=False,
             debug=debug,
             store=store,
+            cache=cache,
             name=name or "LangGraph",
         )
 
@@ -811,6 +820,7 @@ class CompiledStateGraph(CompiledGraph):
                 writers=[ChannelWrite(write_entries)],
                 metadata=node.metadata,
                 retry_policy=node.retry_policy,
+                cache_policy=node.cache_policy,
                 bound=node.runnable,
             )
         else:

libs/langgraph/langgraph/pregel/__init__.py

@@ -36,6 +36,7 @@ from langchain_core.runnables.utils import (
 from pydantic import BaseModel
 from typing_extensions import Self
 
+from langgraph.cache.base import BaseCache
 from langgraph.channels.base import (
     BaseChannel,
 )
@@ -46,7 +47,9 @@ from langgraph.checkpoint.base import (
     copy_checkpoint,
 )
 from langgraph.constants import (
+    CACHE_NS_WRITES,
     CONF,
+    CONFIG_KEY_CACHE,
     CONFIG_KEY_CHECKPOINT_DURING,
     CONFIG_KEY_CHECKPOINT_ID,
     CONFIG_KEY_CHECKPOINT_NS,
@@ -85,6 +88,7 @@ from langgraph.pregel.algo import (
     local_write,
     prepare_next_tasks,
 )
+from langgraph.pregel.call import identifier
 from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint
 from langgraph.pregel.debug import tasks_w_writes
 from langgraph.pregel.draw import draw_graph
@@ -102,6 +106,7 @@ from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
+    CachePolicy,
     Checkpointer,
     Interrupt,
     LoopProtocol,
@@ -495,8 +500,15 @@ class Pregel(PregelProtocol):
     store: BaseStore | None = None
     """Memory store to use for SharedValues. Defaults to None."""
 
-    retry_policy: Sequence[RetryPolicy] | None = None
-    """Retry policies to use when running tasks. Set to None to disable."""
+    cache: BaseCache | None = None
+    """Cache to use for storing node results. Defaults to None."""
+
+    retry_policy: Sequence[RetryPolicy] = ()
+    """Retry policies to use when running tasks. Empty set disables retries."""
+
+    cache_policy: CachePolicy | None = None
+    """Cache policy to use for all nodes. Can be overridden by individual nodes.
+    Defaults to None."""
 
     config_type: type[Any] | None = None
 
@@ -525,7 +537,9 @@ class Pregel(PregelProtocol):
         debug: bool | None = None,
         checkpointer: BaseCheckpointSaver | None = None,
         store: BaseStore | None = None,
-        retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+        cache: BaseCache | None = None,
+        retry_policy: RetryPolicy | Sequence[RetryPolicy] = (),
+        cache_policy: CachePolicy | None = None,
         config_type: type[Any] | None = None,
         input_model: type[BaseModel] | None = None,
         config: RunnableConfig | None = None,
@@ -545,10 +559,11 @@ class Pregel(PregelProtocol):
         self.debug = debug if debug is not None else get_debug()
         self.checkpointer = checkpointer
         self.store = store
-        if isinstance(retry_policy, RetryPolicy):
-            self.retry_policy: Sequence[RetryPolicy] = (retry_policy,)
-        else:
-            self.retry_policy = retry_policy
+        self.cache = cache
+        self.retry_policy = (
+            (retry_policy,) if isinstance(retry_policy, RetryPolicy) else retry_policy
+        )
+        self.cache_policy = cache_policy
         self.config_type = config_type
         self.input_model = input_model
         self.config = config
@@ -2193,6 +2208,7 @@ class Pregel(PregelProtocol):
         All | Sequence[str],
         BaseCheckpointSaver | None,
         BaseStore | None,
+        BaseCache | None,
     ]:
         if config["recursion_limit"] < 1:
             raise ValueError("recursion_limit must be at least 1")
@@ -2225,6 +2241,10 @@ class Pregel(PregelProtocol):
             store: BaseStore | None = config[CONF][CONFIG_KEY_STORE]
         else:
             store = self.store
+        if CONFIG_KEY_CACHE in config.get(CONF, {}):
+            cache: BaseCache | None = config[CONF][CONFIG_KEY_CACHE]
+        else:
+            cache = self.cache
         return (
             debug,
             set(stream_mode),
@@ -2233,6 +2253,7 @@ class Pregel(PregelProtocol):
             interrupt_after,
             checkpointer,
             store,
+            cache,
         )
 
     def stream(
@@ -2405,6 +2426,7 @@ class Pregel(PregelProtocol):
                 interrupt_after_,
                 checkpointer,
                 store,
+                cache,
             ) = self._defaults(
                 config,
                 stream_mode=stream_mode,
@@ -2436,6 +2458,7 @@ class Pregel(PregelProtocol):
                 stream=StreamProtocol(stream.put, stream_modes),
                 config=config,
                 store=store,
+                cache=cache,
                 checkpointer=checkpointer,
                 nodes=self.nodes,
                 specs=self.channels,
@@ -2450,6 +2473,8 @@ class Pregel(PregelProtocol):
                 else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
                 trigger_to_nodes=self.trigger_to_nodes,
                 migrate_checkpoint=self._migrate_checkpoint,
+                retry_policy=self.retry_policy,
+                cache_policy=self.cache_policy,
             ) as loop:
                 # create runner
                 runner = PregelRunner(
@@ -2494,11 +2519,13 @@ class Pregel(PregelProtocol):
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps.
                 while loop.tick(input_keys=self.input_channels):
+                    for task in loop.match_cached_writes():
+                        loop.output_writes(task.id, task.writes, cached=True)
                     for _ in runner.tick(
-                        loop.tasks.values(),
+                        [t for t in loop.tasks.values() if not t.writes],
                         timeout=self.step_timeout,
-                        retry_policy=self.retry_policy,
                         get_waiter=get_waiter,
+                        match_cached_writes=loop.match_cached_writes,
                     ):
                         # emit output
                         yield from output()
@@ -2710,6 +2737,7 @@ class Pregel(PregelProtocol):
                 interrupt_after_,
                 checkpointer,
                 store,
+                cache,
             ) = self._defaults(
                 config,
                 stream_mode=stream_mode,
@@ -2743,6 +2771,7 @@ class Pregel(PregelProtocol):
                 stream=StreamProtocol(stream.put_nowait, stream_modes),
                 config=config,
                 store=store,
+                cache=cache,
                 checkpointer=checkpointer,
                 nodes=self.nodes,
                 specs=self.channels,
@@ -2757,6 +2786,8 @@ class Pregel(PregelProtocol):
                 else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
                 trigger_to_nodes=self.trigger_to_nodes,
                 migrate_checkpoint=self._migrate_checkpoint,
+                retry_policy=self.retry_policy,
+                cache_policy=self.cache_policy,
             ) as loop:
                 # create runner
                 runner = PregelRunner(
@@ -2792,11 +2823,13 @@ class Pregel(PregelProtocol):
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
                 while loop.tick(input_keys=self.input_channels):
+                    for task in await loop.amatch_cached_writes():
+                        loop.output_writes(task.id, task.writes, cached=True)
                     async for _ in runner.atick(
-                        loop.tasks.values(),
+                        [t for t in loop.tasks.values() if not t.writes],
                         timeout=self.step_timeout,
-                        retry_policy=self.retry_policy,
                         get_waiter=get_waiter,
+                        match_cached_writes=loop.amatch_cached_writes,
                     ):
                         # emit output
                         for o in output():
@@ -2958,6 +2991,44 @@ class Pregel(PregelProtocol):
         else:
             return chunks
 
+    def clear_cache(self, nodes: Sequence[str] | None = None) -> None:
+        """Clear the cache for the given nodes."""
+        if not self.cache:
+            raise ValueError("No cache is set for this graph. Cannot clear cache.")
+        nodes = nodes or self.nodes.keys()
+        # collect namespaces to clear
+        namespaces: list[tuple[str, ...]] = []
+        for node in nodes:
+            if node in self.nodes:
+                namespaces.append(
+                    (
+                        CACHE_NS_WRITES,
+                        (identifier(self.nodes[node]) or "__dynamic__"),
+                        node,
+                    ),
+                )
+        # clear cache
+        self.cache.clear(namespaces)
+
+    async def aclear_cache(self, nodes: Sequence[str] | None = None) -> None:
+        """Asynchronously clear the cache for the given nodes."""
+        if not self.cache:
+            raise ValueError("No cache is set for this graph. Cannot clear cache.")
+        nodes = nodes or self.nodes.keys()
+        # collect namespaces to clear
+        namespaces: list[tuple[str, ...]] = []
+        for node in nodes:
+            if node in self.nodes:
+                namespaces.append(
+                    (
+                        CACHE_NS_WRITES,
+                        (identifier(self.nodes[node]) or "__dynamic__"),
+                        node,
+                    ),
+                )
+        # clear cache
+        await self.cache.aclear(namespaces)
+
 
 def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]:
     """Index from a trigger to nodes that depend on it."""

libs/langgraph/langgraph/pregel/algo.py

@@ -33,6 +33,7 @@ from langgraph.checkpoint.base import (
     V,
 )
 from langgraph.constants import (
+    CACHE_NS_WRITES,
     CONF,
     CONFIG_KEY_CHECKPOINT_ID,
     CONFIG_KEY_CHECKPOINT_MAP,
@@ -65,13 +66,15 @@ from langgraph.constants import (
 )
 from langgraph.errors import InvalidUpdateError
 from langgraph.managed.base import ManagedValueMapping
-from langgraph.pregel.call import get_runnable_for_task
+from langgraph.pregel.call import get_runnable_for_task, identifier
 from langgraph.pregel.io import read_channels
 from langgraph.pregel.log import logger
 from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
+    CacheKey,
+    CachePolicy,
     PregelExecutableTask,
     PregelScratchpad,
     PregelTask,
@@ -111,24 +114,27 @@ class PregelTaskWrites(NamedTuple):
 
 
 class Call:
-    __slots__ = ("func", "input", "retry", "callbacks")
+    __slots__ = ("func", "input", "retry", "cache_policy", "callbacks")
 
     func: Callable
-    input: Any
+    input: tuple[tuple[Any, ...], dict[str, Any]]
     retry: Optional[Sequence[RetryPolicy]]
+    cache_policy: Optional[CachePolicy]
     callbacks: Callbacks
 
     def __init__(
         self,
         func: Callable,
-        input: Any,
+        input: tuple[tuple[Any, ...], dict[str, Any]],
         *,
         retry: Optional[Sequence[RetryPolicy]],
+        cache_policy: Optional[CachePolicy],
         callbacks: Callbacks,
     ) -> None:
         self.func = func
         self.input = input
         self.retry = retry
+        self.cache_policy = cache_policy
         self.callbacks = callbacks
 
 
@@ -382,6 +388,8 @@ def prepare_next_tasks(
     manager: Literal[None] = None,
     trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
     updated_channels: Optional[set[str]] = None,
+    retry_policy: Sequence[RetryPolicy] = (),
+    cache_policy: Literal[None] = None,
 ) -> dict[str, PregelTask]: ...
 
 
@@ -401,6 +409,8 @@ def prepare_next_tasks(
     manager: Union[None, ParentRunManager, AsyncParentRunManager],
     trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
     updated_channels: Optional[set[str]] = None,
+    retry_policy: Sequence[RetryPolicy] = (),
+    cache_policy: Optional[CachePolicy] = None,
 ) -> dict[str, PregelExecutableTask]: ...
 
 
@@ -419,6 +429,8 @@ def prepare_next_tasks(
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
     trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
     updated_channels: Optional[set[str]] = None,
+    retry_policy: Sequence[RetryPolicy] = (),
+    cache_policy: Optional[CachePolicy] = None,
 ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
     """Prepare the set of tasks that will make up the next Pregel step.
 
@@ -469,6 +481,8 @@ def prepare_next_tasks(
             checkpointer=checkpointer,
             manager=manager,
             input_cache=input_cache,
+            cache_policy=cache_policy,
+            retry_policy=retry_policy,
         ):
             tasks.append(task)
 
@@ -512,6 +526,8 @@ def prepare_next_tasks(
             checkpointer=checkpointer,
             manager=manager,
             input_cache=input_cache,
+            cache_policy=cache_policy,
+            retry_policy=retry_policy,
         ):
             tasks.append(task)
     return {t.id: t for t in tasks}
@@ -538,6 +554,8 @@ def prepare_single_task(
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
     input_cache: Optional[dict[INPUT_CACHE_KEY_TYPE, Any]] = None,
+    cache_policy: Optional[CachePolicy] = None,
+    retry_policy: Sequence[RetryPolicy] = (),
 ) -> Union[None, PregelTask, PregelExecutableTask]:
     """Prepares a single task for the next Pregel step, given a task path, which
     uniquely identifies a PUSH or PULL task within the graph."""
@@ -580,6 +598,21 @@ def prepare_single_task(
             assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
         if for_execution:
             writes: deque[tuple[str, Any]] = deque()
+            cache_policy = call.cache_policy or cache_policy
+            if cache_policy:
+                args_key = cache_policy.key_func(*call.input[0], **call.input[1])
+                cache_key: Optional[CacheKey] = CacheKey(
+                    (
+                        CACHE_NS_WRITES,
+                        (identifier(call.func) or "__dynamic__"),
+                    ),
+                    xxh3_128_hexdigest(
+                        args_key.encode() if isinstance(args_key, str) else args_key,
+                    ),
+                    cache_policy.ttl,
+                )
+            else:
+                cache_key = None
             return PregelExecutableTask(
                 name,
                 call.input,
@@ -624,8 +657,8 @@ def prepare_single_task(
                     },
                 ),
                 triggers,
-                call.retry,
-                None,
+                call.retry or retry_policy,
+                cache_key,
                 task_id,
                 task_path,
             )
@@ -649,6 +682,11 @@ def prepare_single_task(
                     f"Ignoring unknown node name {packet.node} in pending sends"
                 )
                 return
+            # find process
+            proc = processes[packet.node]
+            proc_node = proc.node
+            if proc_node is None:
+                return
             # create task id
             triggers = PUSH_TRIGGER
             checkpoint_ns = (
@@ -679,73 +717,80 @@ def prepare_single_task(
         if task_id_checksum is not None:
             assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
         if for_execution:
-            proc = processes[packet.node]
-            if node := proc.node:
-                if proc.metadata:
-                    metadata.update(proc.metadata)
-                writes = deque()
-                return PregelExecutableTask(
-                    packet.node,
-                    packet.arg,
-                    node,
-                    writes,
-                    patch_config(
-                        merge_configs(
-                            config, {"metadata": metadata, "tags": proc.tags}
-                        ),
-                        run_name=packet.node,
-                        callbacks=(
-                            manager.get_child(f"graph:step:{step}") if manager else None
-                        ),
-                        configurable={
-                            CONFIG_KEY_TASK_ID: task_id,
-                            # deque.extend is thread-safe
-                            CONFIG_KEY_SEND: partial(
-                                local_write,
-                                writes.extend,
-                                processes.keys(),
-                            ),
-                            CONFIG_KEY_READ: partial(
-                                local_read,
-                                channels,
-                                managed,
-                                PregelTaskWrites(
-                                    task_path, packet.node, writes, triggers
-                                ),
-                            ),
-                            CONFIG_KEY_STORE: (
-                                store or configurable.get(CONFIG_KEY_STORE)
-                            ),
-                            CONFIG_KEY_CHECKPOINTER: (
-                                checkpointer
-                                or configurable.get(CONFIG_KEY_CHECKPOINTER)
-                            ),
-                            CONFIG_KEY_CHECKPOINT_MAP: {
-                                **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
-                                parent_ns: checkpoint["id"],
-                            },
-                            CONFIG_KEY_CHECKPOINT_ID: None,
-                            CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                            CONFIG_KEY_SCRATCHPAD: _scratchpad(
-                                config[CONF].get(CONFIG_KEY_SCRATCHPAD),
-                                pending_writes,
-                                task_id,
-                                xxh3_128_hexdigest(task_checkpoint_ns.encode()),
-                                config[CONF].get(CONFIG_KEY_RESUME_MAP),
-                            ),
-                            CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get(
-                                PREVIOUS, None
-                            ),
-                        },
+            if proc.metadata:
+                metadata.update(proc.metadata)
+            writes = deque()
+            cache_policy = proc.cache_policy or cache_policy
+            if cache_policy:
+                args_key = cache_policy.key_func(packet.arg)
+                cache_key = CacheKey(
+                    (
+                        CACHE_NS_WRITES,
+                        (identifier(proc) or "__dynamic__"),
+                        packet.node,
                     ),
-                    triggers,
-                    proc.retry_policy,
-                    None,
-                    task_id,
-                    task_path,
-                    writers=proc.flat_writers,
-                    subgraphs=proc.subgraphs,
+                    xxh3_128_hexdigest(
+                        args_key.encode() if isinstance(args_key, str) else args_key,
+                    ),
+                    cache_policy.ttl,
                 )
+            else:
+                cache_key = None
+            return PregelExecutableTask(
+                packet.node,
+                packet.arg,
+                proc_node,
+                writes,
+                patch_config(
+                    merge_configs(config, {"metadata": metadata, "tags": proc.tags}),
+                    run_name=packet.node,
+                    callbacks=(
+                        manager.get_child(f"graph:step:{step}") if manager else None
+                    ),
+                    configurable={
+                        CONFIG_KEY_TASK_ID: task_id,
+                        # deque.extend is thread-safe
+                        CONFIG_KEY_SEND: partial(
+                            local_write,
+                            writes.extend,
+                            processes.keys(),
+                        ),
+                        CONFIG_KEY_READ: partial(
+                            local_read,
+                            channels,
+                            managed,
+                            PregelTaskWrites(task_path, packet.node, writes, triggers),
+                        ),
+                        CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)),
+                        CONFIG_KEY_CHECKPOINTER: (
+                            checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER)
+                        ),
+                        CONFIG_KEY_CHECKPOINT_MAP: {
+                            **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
+                            parent_ns: checkpoint["id"],
+                        },
+                        CONFIG_KEY_CHECKPOINT_ID: None,
+                        CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
+                        CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                            config[CONF].get(CONFIG_KEY_SCRATCHPAD),
+                            pending_writes,
+                            task_id,
+                            xxh3_128_hexdigest(task_checkpoint_ns.encode()),
+                            config[CONF].get(CONFIG_KEY_RESUME_MAP),
+                        ),
+                        CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get(
+                            PREVIOUS, None
+                        ),
+                    },
+                ),
+                triggers,
+                proc.retry_policy or retry_policy,
+                cache_key,
+                task_id,
+                task_path,
+                writers=proc.flat_writers,
+                subgraphs=proc.subgraphs,
+            )
         else:
             return PregelTask(task_id, packet.node, task_path)
     elif task_path[0] == PULL:
@@ -807,6 +852,24 @@ def prepare_single_task(
                     if proc.metadata:
                         metadata.update(proc.metadata)
                     writes = deque()
+                    cache_policy = proc.cache_policy or cache_policy
+                    if cache_policy:
+                        args_key = cache_policy.key_func(val)
+                        cache_key = CacheKey(
+                            (
+                                CACHE_NS_WRITES,
+                                (identifier(proc) or "__dynamic__"),
+                                name,
+                            ),
+                            xxh3_128_hexdigest(
+                                args_key.encode()
+                                if isinstance(args_key, str)
+                                else args_key,
+                            ),
+                            cache_policy.ttl,
+                        )
+                    else:
+                        cache_key = None
                     return PregelExecutableTask(
                         name,
                         val,
@@ -867,8 +930,8 @@ def prepare_single_task(
                             },
                         ),
                         triggers,
-                        proc.retry_policy,
-                        None,
+                        proc.retry_policy or retry_policy,
+                        cache_key,
                         task_id,
                         task_path[:3],
                         writers=proc.flat_writers,
@@ -1000,7 +1063,8 @@ def _proc_input(
                     val = channels[chan].get()
                     break
             else:
-                val[k] = managed[k]()
+                val = managed[chan]()
+                break
         else:
             return MISSING
     else:
@@ -1019,18 +1083,20 @@ def _proc_input(
     return val
 
 
-def _uuid5_str(namespace: bytes, *parts: str) -> str:
+def _uuid5_str(namespace: bytes, *parts: Union[str, bytes]) -> str:
     """Generate a UUID from the SHA-1 hash of a namespace and str parts."""
 
     sha = sha1(namespace, usedforsecurity=False)
-    sha.update(b"".join(p.encode() for p in parts))
+    sha.update(b"".join(p.encode() if isinstance(p, str) else p for p in parts))
     hex = sha.hexdigest()
     return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}"
 
 
-def _xxhash_str(namespace: bytes, *parts: str) -> str:
+def _xxhash_str(namespace: bytes, *parts: Union[str, bytes]) -> str:
     """Generate a UUID from the XXH3 hash of a namespace and str parts."""
-    hex = xxh3_128_hexdigest(namespace + b"".join(p.encode() for p in parts))
+    hex = xxh3_128_hexdigest(
+        namespace + b"".join(p.encode() if isinstance(p, str) else p for p in parts)
+    )
     return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}"
 
 

libs/langgraph/langgraph/pregel/call.py

@@ -13,7 +13,7 @@ from typing_extensions import ParamSpec
 
 from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
-from langgraph.types import RetryPolicy
+from langgraph.types import CachePolicy, RetryPolicy
 from langgraph.utils.config import get_config
 from langgraph.utils.runnable import (
     RunnableCallable,
@@ -28,6 +28,7 @@ from langgraph.utils.runnable import (
 
 
 def _getattribute(obj: Any, name: str) -> Any:
+    parent = None
     for subpath in name.split("."):
         if subpath == "<locals>":
             raise AttributeError(f"Can't get local attribute {name!r} on {obj!r}")
@@ -73,6 +74,35 @@ def _whichmodule(obj: Any, name: str) -> Optional[str]:
     return None
 
 
+def identifier(obj: Any, name: Optional[str] = None) -> Optional[str]:
+    """Return the module and name of an object."""
+    from langgraph.pregel.read import PregelNode
+    from langgraph.utils.runnable import RunnableCallable, RunnableSeq
+
+    if isinstance(obj, PregelNode):
+        obj = obj.bound
+    if isinstance(obj, RunnableSeq):
+        obj = obj.steps[0]
+    if isinstance(obj, RunnableCallable):
+        obj = obj.func
+    if name is None:
+        name = getattr(obj, "__qualname__", None)
+    if name is None:  # pragma: no cover
+        # This used to be needed for Python 2.7 support but is probably not
+        # needed anymore. However we keep the __name__ introspection in case
+        # users of cloudpickle rely on this old behavior for unknown reasons.
+        name = getattr(obj, "__name__", None)
+    if name is None:
+        return None
+
+    module_name = getattr(obj, "__module__", None)
+    if module_name is None:
+        # In this case, obj.__module__ is None. obj is thus treated as dynamic.
+        return None
+
+    return f"{module_name}.{name}"
+
+
 def _lookup_module_and_qualname(
     obj: Any, name: Optional[str] = None
 ) -> Optional[tuple[types.ModuleType, str]]:
@@ -135,7 +165,7 @@ def _explode_args_trace_inputs(
     return arguments
 
 
-def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq:
+def get_runnable_for_entrypoint(func: Callable[..., Any]) -> Runnable:
     key = (func, False)
     if key in CACHE:
         return CACHE[key]
@@ -160,7 +190,7 @@ def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq:
         return CACHE.setdefault(key, run)
 
 
-def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq:
+def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
     key = (func, True)
     if key in CACHE:
         return CACHE[key]
@@ -222,9 +252,16 @@ def call(
     func: Callable[P, T],
     *args: Any,
     retry: Optional[Sequence[RetryPolicy]] = 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, callbacks=config["callbacks"])
+    fut = impl(
+        func,
+        (args, kwargs),
+        retry=retry,
+        cache_policy=cache_policy,
+        callbacks=config["callbacks"],
+    )
     return fut

libs/langgraph/langgraph/pregel/loop.py

@@ -27,6 +27,7 @@ from langchain_core.runnables import RunnableConfig
 from pydantic import BaseModel
 from typing_extensions import ParamSpec, Self
 
+from langgraph.cache.base import BaseCache
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import (
     EXCLUDED_METADATA_KEYS,
@@ -117,10 +118,12 @@ from langgraph.pregel.utils import get_new_channel_versions, is_xxh3_128_hexdige
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
+    CachePolicy,
     Command,
     LoopProtocol,
     PregelExecutableTask,
     PregelScratchpad,
+    RetryPolicy,
     StreamChunk,
     StreamProtocol,
 )
@@ -133,6 +136,7 @@ INPUT_DONE = object()
 INPUT_RESUMING = object()
 INPUT_SHOULD_VALIDATE = object()
 SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED)
+WritesT = Sequence[tuple[str, Any]]
 
 
 def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
@@ -147,6 +151,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
 class PregelLoop(LoopProtocol):
     input: Optional[Any]
     input_model: Optional[type[BaseModel]]
+    cache: Optional[BaseCache[WritesT]]
     checkpointer: Optional[BaseCheckpointSaver]
     nodes: Mapping[str, PregelNode]
     specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
@@ -159,18 +164,18 @@ class PregelLoop(LoopProtocol):
     interrupt_before: Union[All, Sequence[str]]
     checkpoint_during: bool
     debug: bool
+    retry_policy: Sequence[RetryPolicy]
+    cache_policy: Optional[CachePolicy]
 
     checkpointer_get_next_version: GetNextVersion
-    checkpointer_put_writes: Optional[
-        Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
-    ]
+    checkpointer_put_writes: Optional[Callable[[RunnableConfig, WritesT, str], Any]]
     checkpointer_put_writes_accepts_task_path: bool
     _checkpointer_put_after_previous: Optional[
         Callable[
             [
                 Optional[concurrent.futures.Future],
                 RunnableConfig,
-                Sequence[tuple[str, Any]],
+                Checkpoint,
                 str,
                 ChannelVersions,
             ],
@@ -206,6 +211,7 @@ class PregelLoop(LoopProtocol):
         stream: Optional[StreamProtocol],
         config: RunnableConfig,
         store: Optional[BaseStore],
+        cache: Optional[BaseCache],
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
@@ -218,6 +224,8 @@ class PregelLoop(LoopProtocol):
         input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
+        retry_policy: Sequence[RetryPolicy] = (),
+        cache_policy: Optional[CachePolicy] = None,
         checkpoint_during: bool = True,
     ) -> None:
         super().__init__(
@@ -230,6 +238,7 @@ class PregelLoop(LoopProtocol):
         self.input = input
         self.input_model = input_model
         self.checkpointer = checkpointer
+        self.cache = cache
         self.nodes = nodes
         self.specs = specs
         self.output_keys = output_keys
@@ -244,6 +253,8 @@ class PregelLoop(LoopProtocol):
         )
         self._migrate_checkpoint = migrate_checkpoint
         self.trigger_to_nodes = trigger_to_nodes
+        self.retry_policy = retry_policy
+        self.cache_policy = cache_policy
         self.checkpoint_during = checkpoint_during
         self.debug = debug
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
@@ -299,7 +310,7 @@ class PregelLoop(LoopProtocol):
         )
         self.prev_checkpoint_config = None
 
-    def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
+    def put_writes(self, task_id: str, writes: WritesT) -> None:
         """Put writes for a task, to be read by the next tick."""
         if not writes:
             return
@@ -346,7 +357,7 @@ class PregelLoop(LoopProtocol):
                 )
         # output writes
         if hasattr(self, "tasks"):
-            self._output_writes(task_id, writes)
+            self.output_writes(task_id, writes)
 
     def _put_pending_writes(self) -> None:
         if self.checkpointer_put_writes is None:
@@ -418,6 +429,8 @@ class PregelLoop(LoopProtocol):
                 store=self.store,
                 checkpointer=self.checkpointer,
                 manager=self.manager,
+                retry_policy=self.retry_policy,
+                cache_policy=self.cache_policy,
             ),
         ):
             # don't start if we should interrupt *before* the new task
@@ -547,6 +560,8 @@ class PregelLoop(LoopProtocol):
             checkpointer=self.checkpointer,
             trigger_to_nodes=self.trigger_to_nodes,
             updated_channels=updated_channels,
+            retry_policy=self.retry_policy,
+            cache_policy=self.cache_policy,
         )
         self.to_interrupt = []
 
@@ -609,10 +624,16 @@ class PregelLoop(LoopProtocol):
         # print output for any tasks we applied previous writes to
         for task in self.tasks.values():
             if task.writes:
-                self._output_writes(task.id, task.writes, cached=True)
+                self.output_writes(task.id, task.writes, cached=True)
 
         return True
 
+    def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        raise NotImplementedError
+
+    async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        raise NotImplementedError
+
     # private
 
     def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None:
@@ -912,8 +933,8 @@ class PregelLoop(LoopProtocol):
         for v in values(*args, **kwargs):
             self.stream((self.checkpoint_ns, mode, v))
 
-    def _output_writes(
-        self, task_id: str, writes: Sequence[tuple[str, Any]], *, cached: bool = False
+    def output_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
     ) -> None:
         if task := self.tasks.get(task_id):
             if task.config is not None and TAG_HIDDEN in task.config.get(
@@ -967,6 +988,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         stream: Optional[StreamProtocol],
         config: RunnableConfig,
         store: Optional[BaseStore],
+        cache: Optional[BaseCache],
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
@@ -979,6 +1001,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
+        retry_policy: Sequence[RetryPolicy] = (),
+        cache_policy: Optional[CachePolicy] = None,
         checkpoint_during: bool = True,
     ) -> None:
         super().__init__(
@@ -987,6 +1011,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
             stream=stream,
             config=config,
             checkpointer=checkpointer,
+            cache=cache,
             store=store,
             nodes=nodes,
             specs=specs,
@@ -998,6 +1023,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
             debug=debug,
             migrate_checkpoint=migrate_checkpoint,
             trigger_to_nodes=trigger_to_nodes,
+            retry_policy=retry_policy,
+            cache_policy=cache_policy,
             checkpoint_during=checkpoint_during,
         )
         self.stack = ExitStack()
@@ -1037,6 +1064,39 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
 
         return self.submit(cast(WritableManagedValue, managed_value).update, values)
 
+    def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        if self.cache is None:
+            return ()
+        matched: list[PregelExecutableTask] = []
+        if cached := {
+            (t.cache_key.ns, t.cache_key.key): t
+            for t in self.tasks.values()
+            if t.cache_key and not t.writes
+        }:
+            for key, values in self.cache.get(tuple(cached)).items():
+                task = cached[key]
+                task.writes.extend(values)
+                matched.append(task)
+        return matched
+
+    def put_writes(self, task_id: str, writes: WritesT) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        super().put_writes(task_id, writes)
+        if not writes or self.cache is None or not hasattr(self, "tasks"):
+            return
+        task = self.tasks.get(task_id)
+        if task is None or task.cache_key is None:
+            return
+        self.submit(
+            self.cache.set,
+            {
+                (task.cache_key.ns, task.cache_key.key): (
+                    task.writes,
+                    task.cache_key.ttl,
+                )
+            },
+        )
+
     # context manager
 
     def __enter__(self) -> Self:
@@ -1117,6 +1177,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         stream: Optional[StreamProtocol],
         config: RunnableConfig,
         store: Optional[BaseStore],
+        cache: Optional[BaseCache],
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
@@ -1129,6 +1190,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
+        retry_policy: Sequence[RetryPolicy] = (),
+        cache_policy: Optional[CachePolicy] = None,
         checkpoint_during: bool = True,
     ) -> None:
         super().__init__(
@@ -1137,6 +1200,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
             stream=stream,
             config=config,
             checkpointer=checkpointer,
+            cache=cache,
             store=store,
             nodes=nodes,
             specs=specs,
@@ -1148,6 +1212,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
             debug=debug,
             migrate_checkpoint=migrate_checkpoint,
             trigger_to_nodes=trigger_to_nodes,
+            retry_policy=retry_policy,
+            cache_policy=cache_policy,
             checkpoint_during=checkpoint_during,
         )
         self.stack = AsyncExitStack()
@@ -1187,6 +1253,42 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
 
         return self.submit(cast(WritableManagedValue, managed_value).aupdate, values)
 
+    async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        if self.cache is None:
+            return []
+        matched: list[PregelExecutableTask] = []
+        if cached := {
+            (t.cache_key.ns, t.cache_key.key): t
+            for t in self.tasks.values()
+            if t.cache_key and not t.writes
+        }:
+            for key, values in (await self.cache.aget(tuple(cached))).items():
+                task = cached[key]
+                task.writes.extend(values)
+                matched.append(task)
+        return matched
+
+    def put_writes(self, task_id: str, writes: WritesT) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        super().put_writes(task_id, writes)
+        if not writes or self.cache is None or not hasattr(self, "tasks"):
+            return
+        task = self.tasks.get(task_id)
+        if task is None or task.cache_key is None:
+            return
+        if writes[0][0] in (INTERRUPT, ERROR):
+            # only cache successful tasks
+            return
+        self.submit(
+            self.cache.aset,
+            {
+                (task.cache_key.ns, task.cache_key.key): (
+                    task.writes,
+                    task.cache_key.ttl,
+                )
+            },
+        )
+
     # context manager
 
     async def __aenter__(self) -> Self:

libs/langgraph/langgraph/pregel/read.py

@@ -14,14 +14,15 @@ from langchain_core.runnables import (
     RunnablePassthrough,
     RunnableSerializable,
 )
-from langchain_core.runnables.base import Input, Other, coerce_to_runnable
-from langchain_core.runnables.utils import ConfigurableFieldSpec
+from langchain_core.runnables.base import Other, coerce_to_runnable
+from langchain_core.runnables.utils import ConfigurableFieldSpec, Input
 
 from langgraph.constants import CONF, CONFIG_KEY_READ
 from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.utils import find_subgraph_pregel
 from langgraph.pregel.write import ChannelWrite
+from langgraph.types import CachePolicy
 from langgraph.utils.config import merge_configs
 from langgraph.utils.runnable import RunnableCallable, RunnableSeq
 
@@ -143,6 +144,9 @@ class PregelNode(Runnable):
     retry_policy: Sequence[RetryPolicy] | None
     """The retry policies to use when invoking the node."""
 
+    cache_policy: CachePolicy | None
+    """The cache policy to use when invoking the node."""
+
     tags: Sequence[str] | None
     """Tags to attach to the node for tracing."""
 
@@ -163,6 +167,7 @@ class PregelNode(Runnable):
         metadata: Mapping[str, Any] | None = None,
         bound: Runnable[Any, Any] | None = None,
         retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
+        cache_policy: CachePolicy | None = None,
         subgraphs: Sequence[PregelProtocol] | None = None,
     ) -> None:
         self.channels = channels
@@ -170,8 +175,9 @@ class PregelNode(Runnable):
         self.mapper = mapper
         self.writers = writers or []
         self.bound = bound if bound is not None else DEFAULT_BOUND
+        self.cache_policy = cache_policy
         if isinstance(retry_policy, RetryPolicy):
-            self.retry_policy: Sequence[RetryPolicy] = (retry_policy,)
+            self.retry_policy = (retry_policy,)
         else:
             self.retry_policy = retry_policy
         self.tags = tags

libs/langgraph/langgraph/pregel/retry.py

@@ -3,9 +3,9 @@ import logging
 import random
 import sys
 import time
-from collections.abc import Sequence
+from collections.abc import Awaitable, Sequence
 from dataclasses import replace
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 from langgraph.constants import (
     CONF,
@@ -61,7 +61,7 @@ def run_with_retry(
         except Exception as exc:
             if SUPPORTS_EXC_NOTES:
                 exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
-            if retry_policy is None:
+            if not retry_policy:
                 raise
 
             # Check which retry policy applies to this exception
@@ -106,6 +106,9 @@ async def arun_with_retry(
     task: PregelExecutableTask,
     retry_policies: Optional[Sequence[RetryPolicy]],
     stream: bool = False,
+    match_cached_writes: Optional[
+        Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
+    ] = None,
     configurable: Optional[dict[str, Any]] = None,
 ) -> None:
     """Run a task asynchronously with retries."""
@@ -114,6 +117,11 @@ async def arun_with_retry(
     config = task.config
     if configurable is not None:
         config = patch_configurable(config, configurable)
+    if match_cached_writes is not None and task.cache_key is not None:
+        for t in await match_cached_writes():
+            if t is task:
+                # if the task is already cached, return
+                return
     while True:
         try:
             # clear any writes from previous attempts
@@ -149,7 +157,7 @@ async def arun_with_retry(
         except Exception as exc:
             if SUPPORTS_EXC_NOTES:
                 exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
-            if retry_policies is None:
+            if not retry_policies:
                 raise
 
             # Check which retry policy applies to this exception

libs/langgraph/langgraph/pregel/runner.py

@@ -33,7 +33,12 @@ from langgraph.errors import GraphBubbleUp, GraphInterrupt
 from langgraph.pregel.algo import Call
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
-from langgraph.types import PregelExecutableTask, PregelScratchpad, RetryPolicy
+from langgraph.types import (
+    CachePolicy,
+    PregelExecutableTask,
+    PregelScratchpad,
+    RetryPolicy,
+)
 from langgraph.utils.future import chain_future
 
 F = TypeVar("F", concurrent.futures.Future, asyncio.Future)
@@ -137,6 +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[[], Sequence[PregelExecutableTask]]
+        ] = None,
     ) -> Iterator[None]:
         tasks = tuple(tasks)
         futures = FuturesDict(
@@ -147,7 +155,9 @@ class PregelRunner:
         # give control back to the caller
         yield
         # fast path if single task with no timeout and no waiter
-        if len(tasks) == 1 and timeout is None and get_waiter is None:
+        if len(tasks) == 0:
+            return
+        elif len(tasks) == 1 and timeout is None and get_waiter is None:
             t = tasks[0]
             try:
                 run_with_retry(
@@ -160,6 +170,7 @@ class PregelRunner:
                             retry=retry_policy,
                             futures=weakref.ref(futures),
                             schedule_task=self.schedule_task,
+                            match_cached_writes=match_cached_writes,
                             submit=self.submit,
                             reraise=reraise,
                         ),
@@ -191,25 +202,25 @@ class PregelRunner:
             futures[get_waiter()] = None
         # schedule tasks
         for t in tasks:
-            if not t.writes:
-                fut = self.submit()(  # type: ignore[misc]
-                    run_with_retry,
-                    t,
-                    retry_policy,
-                    configurable={
-                        CONFIG_KEY_CALL: partial(
-                            _call,
-                            weakref.ref(t),
-                            retry=retry_policy,
-                            futures=weakref.ref(futures),
-                            schedule_task=self.schedule_task,
-                            submit=self.submit,
-                            reraise=reraise,
-                        ),
-                    },
-                    __reraise_on_exit__=reraise,
-                )
-                futures[fut] = t
+            fut = self.submit()(  # type: ignore[misc]
+                run_with_retry,
+                t,
+                retry_policy,
+                configurable={
+                    CONFIG_KEY_CALL: partial(
+                        _call,
+                        weakref.ref(t),
+                        retry=retry_policy,
+                        futures=weakref.ref(futures),
+                        schedule_task=self.schedule_task,
+                        match_cached_writes=match_cached_writes,
+                        submit=self.submit,
+                        reraise=reraise,
+                    ),
+                },
+                __reraise_on_exit__=reraise,
+            )
+            futures[fut] = t
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
@@ -266,6 +277,9 @@ class PregelRunner:
         timeout: Optional[float] = None,
         retry_policy: Optional[Sequence[RetryPolicy]] = None,
         get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
+        match_cached_writes: Optional[
+            Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
+        ] = None,
     ) -> AsyncIterator[None]:
         loop = asyncio.get_event_loop()
         tasks = tuple(tasks)
@@ -277,7 +291,9 @@ class PregelRunner:
         # give control back to the caller
         yield
         # fast path if single task with no waiter and no timeout
-        if len(tasks) == 1 and get_waiter is None and timeout is None:
+        if len(tasks) == 0:
+            return
+        elif len(tasks) == 1 and get_waiter is None and timeout is None:
             t = tasks[0]
             try:
                 await arun_with_retry(
@@ -292,6 +308,7 @@ class PregelRunner:
                             retry=retry_policy,
                             futures=weakref.ref(futures),
                             schedule_task=self.schedule_task,
+                            match_cached_writes=match_cached_writes,
                             submit=self.submit,
                             reraise=reraise,
                             loop=loop,
@@ -324,33 +341,33 @@ class PregelRunner:
             futures[get_waiter()] = None
         # schedule tasks
         for t in tasks:
-            if not t.writes:
-                fut = cast(
-                    asyncio.Future,
-                    self.submit()(  # type: ignore[misc]
-                        arun_with_retry,
-                        t,
-                        retry_policy,
-                        stream=self.use_astream,
-                        configurable={
-                            CONFIG_KEY_CALL: partial(
-                                _acall,
-                                weakref.ref(t),
-                                retry=retry_policy,
-                                stream=self.use_astream,
-                                futures=weakref.ref(futures),
-                                schedule_task=self.schedule_task,
-                                submit=self.submit,
-                                reraise=reraise,
-                                loop=loop,
-                            ),
-                        },
-                        __name__=t.name,
-                        __cancel_on_exit__=True,
-                        __reraise_on_exit__=reraise,
-                    ),
-                )
-                futures[fut] = t
+            fut = cast(
+                asyncio.Future,
+                self.submit()(  # type: ignore[misc]
+                    arun_with_retry,
+                    t,
+                    retry_policy,
+                    stream=self.use_astream,
+                    configurable={
+                        CONFIG_KEY_CALL: partial(
+                            _acall,
+                            weakref.ref(t),
+                            retry=retry_policy,
+                            stream=self.use_astream,
+                            futures=weakref.ref(futures),
+                            schedule_task=self.schedule_task,
+                            match_cached_writes=match_cached_writes,
+                            submit=self.submit,
+                            reraise=reraise,
+                            loop=loop,
+                        ),
+                    },
+                    __name__=t.name,
+                    __cancel_on_exit__=True,
+                    __reraise_on_exit__=reraise,
+                ),
+            )
+            futures[fut] = t
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
@@ -515,6 +532,7 @@ def _call(
     input: Any,
     *,
     retry: Optional[Sequence[RetryPolicy]] = None,
+    cache_policy: Optional[CachePolicy] = None,
     callbacks: Callbacks = None,
     futures: weakref.ref[FuturesDict],
     schedule_task: weakref.ref[
@@ -522,6 +540,7 @@ def _call(
             [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
         ]
     ],
+    match_cached_writes: Optional[Callable[[], Sequence[PregelExecutableTask]]],
     submit: weakref.ref[Submit],
     reraise: bool,
 ) -> concurrent.futures.Future[Any]:
@@ -535,8 +554,10 @@ def _call(
     if next_task := schedule_task()(  # type: ignore[misc]
         task(),  # type: ignore[arg-type]
         scratchpad.call_counter(),
-        Call(func, input, retry=retry, callbacks=callbacks),
+        Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks),
     ):
+        if match_cached_writes:
+            match_cached_writes()
         if fut := next(
             (
                 f
@@ -574,6 +595,7 @@ def _call(
                         retry=retry,
                         callbacks=callbacks,
                         schedule_task=schedule_task,
+                        match_cached_writes=match_cached_writes,
                         submit=submit,
                         reraise=reraise,
                     ),
@@ -596,6 +618,7 @@ def _acall(
     input: Any,
     *,
     retry: Optional[Sequence[RetryPolicy]] = None,
+    cache_policy: Optional[CachePolicy] = None,
     callbacks: Callbacks = None,
     # injected dependencies
     futures: weakref.ref[FuturesDict],
@@ -604,6 +627,9 @@ def _acall(
             [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
         ]
     ],
+    match_cached_writes: Optional[
+        Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
+    ] = None,
     submit: weakref.ref[Submit],
     loop: asyncio.AbstractEventLoop,
     reraise: bool = False,
@@ -616,7 +642,7 @@ def _acall(
     if next_task := schedule_task()(  # type: ignore[misc]
         task(),  # type: ignore[arg-type]
         scratchpad.call_counter(),
-        Call(func, input, retry=retry, callbacks=callbacks),
+        Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks),
     ):
         if fut := next(
             (
@@ -652,6 +678,7 @@ def _acall(
                     next_task,
                     retry,
                     stream=stream,
+                    match_cached_writes=match_cached_writes,
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _acall,
@@ -659,6 +686,7 @@ def _acall(
                             stream=stream,
                             futures=futures,
                             schedule_task=schedule_task,
+                            match_cached_writes=match_cached_writes,
                             submit=submit,
                             loop=loop,
                             reraise=reraise,

libs/langgraph/langgraph/types.py

@@ -22,6 +22,7 @@ from typing_extensions import Self
 from xxhash import xxh3_128_hexdigest
 
 from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
+from langgraph.utils.cache import default_cache_key
 from langgraph.utils.fields import get_update_as_tuples
 
 if TYPE_CHECKING:
@@ -122,13 +123,19 @@ 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."""
 
 
-class CachePolicy(NamedTuple):
-    """Configuration for caching nodes.
+KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., Union[str, bytes]])
 
-    !!! version-added "Added in version 0.2.24."
-    """
 
-    pass
+@dataclasses.dataclass(**_DC_KWARGS)
+class CachePolicy(Generic[KeyFuncT]):
+    """Configuration for caching nodes."""
+
+    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."""
+
+    ttl: Optional[int] = None
+    """Time to live for the cache entry in seconds. If None, the entry never expires."""
 
 
 @dataclasses.dataclass(**_DC_KWARGS)
@@ -174,6 +181,17 @@ else:
     _T_DC_KWARGS = {"frozen": True}
 
 
+class CacheKey(NamedTuple):
+    """Cache key for a task."""
+
+    ns: tuple[str, ...]
+    """Namespace for the cache entry."""
+    key: str
+    """Key for the cache entry."""
+    ttl: Optional[int]
+    """Time to live for the cache entry in seconds."""
+
+
 @dataclasses.dataclass(**_T_DC_KWARGS)
 class PregelExecutableTask:
     name: str
@@ -182,8 +200,8 @@ class PregelExecutableTask:
     writes: deque[tuple[str, Any]]
     config: RunnableConfig
     triggers: Sequence[str]
-    retry_policy: Optional[Sequence[RetryPolicy]]
-    cache_policy: Optional[CachePolicy]
+    retry_policy: Sequence[RetryPolicy]
+    cache_key: Optional[CacheKey]
     id: str
     path: tuple[Union[str, int, tuple], ...]
     scheduled: bool = False
@@ -313,7 +331,7 @@ class Command(Generic[N], ToolOutputMixin):
 
     graph: Optional[str] = None
     update: Optional[Any] = None
-    resume: Optional[Union[Any, dict[str, Any]]] = None
+    resume: Optional[Union[dict[str, Any], Any]] = None
     goto: Union[Send, Sequence[Union[Send, N]], N] = ()
 
     def __repr__(self) -> str:

libs/langgraph/langgraph/utils/cache.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from collections.abc import Hashable, Mapping, Sequence
+from typing import Any
+
+
+def _freeze(obj: Any, depth: int = 10) -> Hashable:
+    if isinstance(obj, Hashable) or depth <= 0:
+        # already hashable, no need to freeze
+        return obj
+    elif isinstance(obj, Mapping):
+        # sort keys so {"a":1,"b":2} == {"b":2,"a":1}
+        return tuple(sorted((k, _freeze(v, depth - 1)) for k, v in obj.items()))
+    elif isinstance(obj, Sequence):
+        return tuple(_freeze(x, depth - 1) for x in obj)
+    # numpy / pandas etc. can provide their own .tobytes()
+    elif hasattr(obj, "tobytes"):
+        return (
+            type(obj).__name__,
+            obj.tobytes(),
+            obj.shape if hasattr(obj, "shape") else None,
+        )
+    return obj  # strings, ints, dataclasses with frozen=True, etc.
+
+
+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
+
+    # protocol 5 strikes a good balance between speed and size
+    return pickle.dumps((_freeze(args), _freeze(kwargs)), protocol=5, fix_imports=False)

libs/langgraph/tests/conftest.py

@@ -1,5 +1,5 @@
 import sys
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Iterator
 from contextlib import asynccontextmanager
 from typing import Optional
 from uuid import UUID, uuid4
@@ -11,6 +11,9 @@ from psycopg import AsyncConnection, Connection
 from psycopg_pool import AsyncConnectionPool, ConnectionPool
 from pytest_mock import MockerFixture
 
+from langgraph.cache.base import BaseCache
+from langgraph.cache.memory import InMemoryCache
+from langgraph.cache.sqlite import SqliteCache
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
 from langgraph.checkpoint.postgres.aio import (
@@ -361,6 +364,16 @@ async def _store_postgres_aio_pool():
             await conn.execute(f"DROP DATABASE {database}")
 
 
+@pytest.fixture(scope="function", params=["sqlite", "memory"])
+def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]:
+    if request.param == "sqlite":
+        yield SqliteCache(path=":memory:")
+    elif request.param == "memory":
+        yield InMemoryCache()
+    else:
+        raise ValueError(f"Unknown cache type: {request.param}")
+
+
 @pytest.fixture(scope="function")
 def store_postgres():
     database = f"test_{uuid4().hex[:16]}"

libs/langgraph/tests/test_pregel.py

@@ -14,14 +14,7 @@ from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from random import randrange
-from typing import (
-    Annotated,
-    Any,
-    Literal,
-    Optional,
-    Union,
-    get_type_hints,
-)
+from typing import Annotated, Any, Literal, Optional, Union, get_type_hints
 
 import httpx
 import pytest
@@ -38,6 +31,7 @@ from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
 from typing_extensions import TypedDict
 
+from langgraph.cache.base import BaseCache
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.context import Context
@@ -64,6 +58,7 @@ from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.types import (
+    CachePolicy,
     Command,
     Interrupt,
     PregelTask,
@@ -3571,7 +3566,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
     ]
 
 
-def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
+@pytest.mark.parametrize("with_cache", [True, False])
+def test_in_one_fan_out_state_graph_waiting_edge_multiple(
+    with_cache: bool, cache: BaseCache
+) -> None:
     def sorted_add(
         x: list[str], y: Union[list[str], list[tuple[str, str]]]
     ) -> list[str]:
@@ -3586,7 +3584,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         answer: str
         docs: Annotated[list[str], sorted_add]
 
+    rewrite_query_count = 0
+
     def rewrite_query(data: State) -> State:
+        nonlocal rewrite_query_count
+        rewrite_query_count += 1
         return {"query": f"query: {data['query']}"}
 
     def analyzer_one(data: State) -> State:
@@ -3613,7 +3615,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
 
     workflow = StateGraph(State)
 
-    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node(
+        "rewrite_query",
+        rewrite_query,
+        cache_policy=CachePolicy() if with_cache else None,
+    )
     workflow.add_node("analyzer_one", analyzer_one)
     workflow.add_node("retriever_one", retriever_one)
     workflow.add_node("retriever_two", retriever_two)
@@ -3628,7 +3634,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
     workflow.add_conditional_edges("decider", decider_cond)
     workflow.set_finish_point("qa")
 
-    app = workflow.compile()
+    app = workflow.compile(cache=cache)
 
     assert app.invoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: analyzed: query: what is weather in sf",
@@ -3637,12 +3643,24 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
     }
 
     assert [*app.stream({"query": "what is weather in sf"})] == [
-        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "rewrite_query": {"query": "query: what is weather in sf"},
+            "__metadata__": {"cached": True},
+        }
+        if with_cache
+        else {"rewrite_query": {"query": "query: what is weather in sf"}},
         {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
         {"decider": None},
-        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
+        {
+            "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"},
+            "__metadata__": {"cached": True},
+        }
+        if with_cache
+        else {
+            "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}
+        },
         {
             "analyzer_one": {
                 "query": "analyzed: query: analyzed: query: what is weather in sf"
@@ -3653,6 +3671,18 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         {"decider": None},
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
+    assert rewrite_query_count == 2 if with_cache else 4
+
+    # clear the cache
+    if with_cache:
+        app.clear_cache()
+
+        assert app.invoke({"query": "what is weather in sf"}) == {
+            "query": "analyzed: query: analyzed: query: what is weather in sf",
+            "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+            "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+        }
+        assert rewrite_query_count == 4
 
 
 def test_callable_in_conditional_edges_with_no_path_map() -> None:
@@ -6559,7 +6589,7 @@ def test_falsy_return_from_task(
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_multiple_interrupts_functional(
-    request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
+    request: pytest.FixtureRequest, checkpointer_name: str
 ):
     """Test multiple interrupts with functional API."""
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
@@ -6596,6 +6626,81 @@ def test_multiple_interrupts_functional(
     assert counter == 3
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_multiple_interrupts_functional_cache(
+    request: pytest.FixtureRequest, checkpointer_name: str, cache: BaseCache
+):
+    """Test multiple interrupts with functional API."""
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    counter = 0
+
+    @task(cache_policy=CachePolicy())
+    def double(x: int) -> int:
+        """Increment the counter."""
+        nonlocal counter
+        counter += 1
+        return 2 * x
+
+    @entrypoint(checkpointer=checkpointer, cache=cache)
+    def graph(state: dict) -> dict:
+        """React tool."""
+
+        values = []
+
+        for idx in [1, 1, 2, 2, 3, 3]:
+            values.extend([double(idx).result(), interrupt({"a": "boo"})])
+
+        return {"values": values}
+
+    configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+    graph.invoke({}, configurable)
+    graph.invoke(Command(resume="a"), configurable)
+    graph.invoke(Command(resume="b"), configurable)
+    graph.invoke(Command(resume="c"), configurable)
+    graph.invoke(Command(resume="d"), configurable)
+    graph.invoke(Command(resume="e"), configurable)
+    result = graph.invoke(Command(resume="f"), configurable)
+    # `double` value should be cached appropriately when used w/ `interrupt`
+    assert result == {
+        "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
+    }
+    assert counter == 3
+
+    # should all be cached now
+    configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+    graph.invoke({}, configurable)
+    graph.invoke(Command(resume="a"), configurable)
+    graph.invoke(Command(resume="b"), configurable)
+    graph.invoke(Command(resume="c"), configurable)
+    graph.invoke(Command(resume="d"), configurable)
+    graph.invoke(Command(resume="e"), configurable)
+    result = graph.invoke(Command(resume="f"), configurable)
+    # `double` value should be cached appropriately when used w/ `interrupt`
+    assert result == {
+        "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
+    }
+    assert counter == 3
+
+    # clear cache
+    double.clear_cache(cache)
+
+    # should recompute now
+    configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+    graph.invoke({}, configurable)
+    graph.invoke(Command(resume="a"), configurable)
+    graph.invoke(Command(resume="b"), configurable)
+    graph.invoke(Command(resume="c"), configurable)
+    graph.invoke(Command(resume="d"), configurable)
+    graph.invoke(Command(resume="e"), configurable)
+    result = graph.invoke(Command(resume="f"), configurable)
+    # `double` value should be cached appropriately when used w/ `interrupt`
+    assert result == {
+        "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
+    }
+    assert counter == 6
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_double_interrupt_subgraph(
     request: pytest.FixtureRequest, checkpointer_name: str

libs/langgraph/tests/test_pregel_async.py

@@ -31,6 +31,7 @@ from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
 from typing_extensions import TypedDict
 
+from langgraph.cache.base import BaseCache
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.context import Context
@@ -55,6 +56,7 @@ from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.types import (
+    CachePolicy,
     Command,
     Interrupt,
     PregelTask,
@@ -5395,7 +5397,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
         ]
 
 
-async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
+@pytest.mark.parametrize("with_cache", [True, False])
+async def test_in_one_fan_out_state_graph_waiting_edge_multiple(
+    with_cache: bool, cache: BaseCache
+) -> None:
     def sorted_add(
         x: list[str], y: Union[list[str], list[tuple[str, str]]]
     ) -> list[str]:
@@ -5410,7 +5415,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         answer: str
         docs: Annotated[list[str], sorted_add]
 
+    rewrite_query_count = 0
+
     async def rewrite_query(data: State) -> State:
+        nonlocal rewrite_query_count
+        rewrite_query_count += 1
         return {"query": f"query: {data['query']}"}
 
     async def analyzer_one(data: State) -> State:
@@ -5437,7 +5446,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
 
     workflow = StateGraph(State)
 
-    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node(
+        "rewrite_query",
+        rewrite_query,
+        cache_policy=CachePolicy() if with_cache else None,
+    )
     workflow.add_node("analyzer_one", analyzer_one)
     workflow.add_node("retriever_one", retriever_one)
     workflow.add_node("retriever_two", retriever_two)
@@ -5452,21 +5465,34 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
     workflow.add_conditional_edges("decider", decider_cond)
     workflow.set_finish_point("qa")
 
-    app = workflow.compile()
+    app = workflow.compile(cache=cache)
 
     assert await app.ainvoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: analyzed: query: what is weather in sf",
         "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
         "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
     }
+    assert rewrite_query_count == 2
 
     assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
-        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "rewrite_query": {"query": "query: what is weather in sf"},
+            "__metadata__": {"cached": True},
+        }
+        if with_cache
+        else {"rewrite_query": {"query": "query: what is weather in sf"}},
         {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
         {"decider": None},
-        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
+        {
+            "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"},
+            "__metadata__": {"cached": True},
+        }
+        if with_cache
+        else {
+            "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}
+        },
         {
             "analyzer_one": {
                 "query": "analyzed: query: analyzed: query: what is weather in sf"
@@ -5477,6 +5503,18 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         {"decider": None},
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
+    assert rewrite_query_count == 2 if with_cache else 4
+
+    # clear the cache
+    if with_cache:
+        await app.aclear_cache()
+
+        assert await app.ainvoke({"query": "what is weather in sf"}) == {
+            "query": "analyzed: query: analyzed: query: what is weather in sf",
+            "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+            "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+        }
+        assert rewrite_query_count == 4
 
 
 async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
@@ -7510,6 +7548,80 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None:
         assert counter == 3
 
 
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_multiple_interrupts_functional_cache(
+    checkpointer_name: str, cache: BaseCache
+):
+    """Test multiple interrupts with functional API."""
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        counter = 0
+
+        @task(cache_policy=CachePolicy())
+        def double(x: int) -> int:
+            """Increment the counter."""
+            nonlocal counter
+            counter += 1
+            return 2 * x
+
+        @entrypoint(checkpointer=checkpointer, cache=cache)
+        def graph(state: dict) -> dict:
+            """React tool."""
+
+            values = []
+
+            for idx in [1, 1, 2, 2, 3, 3]:
+                values.extend([double(idx).result(), interrupt({"a": "boo"})])
+
+            return {"values": values}
+
+        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+        await graph.ainvoke({}, configurable)
+        await graph.ainvoke(Command(resume="a"), configurable)
+        await graph.ainvoke(Command(resume="b"), configurable)
+        await graph.ainvoke(Command(resume="c"), configurable)
+        await graph.ainvoke(Command(resume="d"), configurable)
+        await graph.ainvoke(Command(resume="e"), configurable)
+        result = await graph.ainvoke(Command(resume="f"), configurable)
+        # `double` value should be cached appropriately when used w/ `interrupt`
+        assert result == {
+            "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
+        }
+        assert counter == 3
+
+        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+        await graph.ainvoke({}, configurable)
+        await graph.ainvoke(Command(resume="a"), configurable)
+        await graph.ainvoke(Command(resume="b"), configurable)
+        await graph.ainvoke(Command(resume="c"), configurable)
+        await graph.ainvoke(Command(resume="d"), configurable)
+        await graph.ainvoke(Command(resume="e"), configurable)
+        result = await graph.ainvoke(Command(resume="f"), configurable)
+        # `double` value should be cached appropriately when used w/ `interrupt`
+        assert result == {
+            "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
+        }
+        assert counter == 3
+
+        # clear the cache
+        await double.aclear_cache(cache)
+
+        # now should recompute
+        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+        await graph.ainvoke({}, configurable)
+        await graph.ainvoke(Command(resume="a"), configurable)
+        await graph.ainvoke(Command(resume="b"), configurable)
+        await graph.ainvoke(Command(resume="c"), configurable)
+        await graph.ainvoke(Command(resume="d"), configurable)
+        await graph.ainvoke(Command(resume="e"), configurable)
+        result = await graph.ainvoke(Command(resume="f"), configurable)
+        # `double` value should be cached appropriately when used w/ `interrupt`
+        assert result == {
+            "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"],
+        }
+        assert counter == 6
+
+
 @NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_double_interrupt_subgraph(checkpointer_name: str) -> None:

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -151,6 +151,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
             msg["input"],
             config=ensure_config(msg["config"]),
             stream=None,
+            cache=self.graph.cache,
             store=self.graph.store,
             checkpointer=self.graph.checkpointer,
             nodes=graph.nodes,
@@ -338,6 +339,7 @@ class KafkaOrchestrator(AbstractContextManager):
             msg["input"],
             config=ensure_config(msg["config"]),
             stream=None,
+            cache=self.graph.cache,
             store=self.graph.store,
             checkpointer=self.graph.checkpointer,
             nodes=graph.nodes,