WIP

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

@@ -254,6 +254,23 @@ class BaseCheckpointSaver(Generic[V]):
         """
         raise NotImplementedError
 
+    def get_writes(
+        self,
+        task_ids: Sequence[str],
+    ) -> dict[str, list[tuple[str, Any]]]:
+        """Fetch writes associated with the given task IDs.
+
+        Args:
+            task_ids (Sequence[str]): List of task IDs to fetch writes for.
+
+        Returns:
+            dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
+        raise NotImplementedError
+
     def list(
         self,
         config: Optional[RunnableConfig],
@@ -358,6 +375,23 @@ class BaseCheckpointSaver(Generic[V]):
         """
         raise NotImplementedError
 
+    async def aget_writes(
+        self,
+        task_ids: Sequence[str],
+    ) -> dict[str, list[tuple[str, Any]]]:
+        """Asynchronously fetch writes associated with the given task IDs.
+
+        Args:
+            task_ids (Sequence[str]): List of task IDs to fetch writes for.
+
+        Returns:
+            dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes.
+
+        Raises:
+            NotImplementedError: Implement this method in your custom checkpoint saver.
+        """
+        raise NotImplementedError
+
     async def alist(
         self,
         config: Optional[RunnableConfig],

libs/langgraph/langgraph/constants.py

@@ -106,6 +106,7 @@ NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000")
 # the task_id to use for writes that are not associated with a task
 CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
 # holds a mapping of task ns -> resume value for resuming tasks
+TASK_CACHE_NAMESPACE = b"__pregel_task_cache"
 
 RESERVED = {
     TAG_HIDDEN,

libs/langgraph/langgraph/func/__init__.py

@@ -31,7 +31,7 @@ from langgraph.pregel.call import (
 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
 
 
 @overload
@@ -39,6 +39,7 @@ def task(
     *,
     name: Optional[str] = None,
     retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+    cache: Optional[CachePolicy[P]] = None,
 ) -> Callable[
     [Union[Callable[P, Awaitable[T]], Callable[P, T]]],
     Callable[P, SyncAsyncFuture[T]],
@@ -56,6 +57,7 @@ def task(
     *,
     name: Optional[str] = None,
     retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
+    cache: Optional[CachePolicy[P]] = None,
 ) -> Union[
     Callable[
         [Union[Callable[P, Awaitable[T]], Callable[P, T]]],
@@ -141,7 +143,7 @@ def task(
                 # handle regular functions / partials / callable classes, etc.
                 func.__name__ = name
 
-        call_func = functools.partial(call, func, retry=retry_policies)
+        call_func = functools.partial(call, func, retry=retry_policies, cache=cache)
         object.__setattr__(call_func, "_is_pregel_task", True)
         return functools.update_wrapper(call_func, func)
 

libs/langgraph/langgraph/pregel/algo.py

@@ -60,6 +60,7 @@ from langgraph.constants import (
     RESUME,
     RETURN,
     TAG_HIDDEN,
+    TASK_CACHE_NAMESPACE,
     TASKS,
     Send,
 )
@@ -72,6 +73,7 @@ from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
+    CachePolicy,
     PregelExecutableTask,
     PregelScratchpad,
     PregelTask,
@@ -111,24 +113,27 @@ class PregelTaskWrites(NamedTuple):
 
 
 class Call:
-    __slots__ = ("func", "input", "retry", "callbacks")
+    __slots__ = ("func", "input", "retry", "cache", "callbacks")
 
     func: Callable
-    input: Any
+    input: tuple[tuple[Any, ...], dict[str, Any]]
     retry: Optional[Sequence[RetryPolicy]]
+    cache: Optional[CachePolicy]
     callbacks: Callbacks
 
     def __init__(
         self,
         func: Callable,
-        input: Any,
+        input: tuple[tuple[Any, ...], dict[str, Any]],
         *,
         retry: Optional[Sequence[RetryPolicy]],
+        cache: Optional[CachePolicy],
         callbacks: Callbacks,
     ) -> None:
         self.func = func
         self.input = input
         self.retry = retry
+        self.cache = cache
         self.callbacks = callbacks
 
 
@@ -556,15 +561,23 @@ def prepare_single_task(
         # create task id
         triggers: Sequence[str] = PUSH_TRIGGER
         checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
-        task_id = task_id_func(
-            checkpoint_id_bytes,
-            checkpoint_ns,
-            str(step),
-            name,
-            PUSH,
-            task_path_str(task_path[1]),
-            str(task_path[2]),
-        )
+        if call.cache:
+            task_id = task_id_func(
+                TASK_CACHE_NAMESPACE,
+                name,
+                call.cache.key(*call.input[0], **call.input[1]),
+                # TODO add truncated iso timestamp here for ttl
+            )
+        else:
+            task_id = task_id_func(
+                checkpoint_id_bytes,
+                checkpoint_ns,
+                str(step),
+                name,
+                PUSH,
+                task_path_str(task_path[1]),
+                str(task_path[2]),
+            )
         task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
         # we append True to the task path to indicate that a call is being
         # made, so we should not return interrupts from this task (responsibility lies with the parent)
@@ -625,7 +638,7 @@ def prepare_single_task(
                 ),
                 triggers,
                 call.retry,
-                None,
+                call.cache,
                 task_id,
                 task_path,
             )
@@ -649,19 +662,31 @@ 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 = (
                 f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
             )
-            task_id = task_id_func(
-                checkpoint_id_bytes,
-                checkpoint_ns,
-                str(step),
-                packet.node,
-                PUSH,
-                str(idx),
-            )
+            if proc.cache_policy:
+                task_id = task_id_func(
+                    TASK_CACHE_NAMESPACE,
+                    packet.node,
+                    proc.cache_policy.key(packet.arg),
+                )
+            else:
+                task_id = task_id_func(
+                    checkpoint_id_bytes,
+                    checkpoint_ns,
+                    str(step),
+                    packet.node,
+                    PUSH,
+                    str(idx),
+                )
         else:
             logger.warning(f"Ignoring invalid PUSH task path {task_path}")
             return
@@ -679,73 +704,64 @@ 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()
+            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
                     ),
-                    triggers,
-                    proc.retry_policy,
-                    None,
-                    task_id,
-                    task_path,
-                    writers=proc.flat_writers,
-                    subgraphs=proc.subgraphs,
-                )
+                    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,
+                proc.cache_policy,
+                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:
@@ -784,6 +800,7 @@ def prepare_single_task(
 
             # create task id
             checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
+            # TODO implement cache
             task_id = task_id_func(
                 checkpoint_id_bytes,
                 checkpoint_ns,
@@ -1000,7 +1017,8 @@ def _proc_input(
                     val = channels[chan].get()
                     break
             else:
-                val[k] = managed[k]()
+                val = managed[chan]()
+                break
         else:
             return MISSING
     else:
@@ -1019,18 +1037,20 @@ def _proc_input(
     return val
 
 
-def _uuid5_str(namespace: bytes, *parts: str) -> str:
+def _uuid5_str(namespace: bytes, *parts: 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: 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}")
@@ -135,7 +136,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 +161,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 +223,12 @@ def call(
     func: Callable[P, T],
     *args: Any,
     retry: Optional[Sequence[RetryPolicy]] = None,
+    cache: 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=cache, callbacks=config["callbacks"]
+    )
     return fut

libs/langgraph/langgraph/pregel/loop.py

@@ -434,7 +434,10 @@ class PregelLoop(LoopProtocol):
             # save the new task
             self.tasks[pushed.id] = pushed
             # match any pending writes to the new task
-            if self.skip_done_tasks:
+            if call.cache:
+                pass
+                # how to call async method here...
+            elif self.skip_done_tasks:
                 self._match_writes({pushed.id: pushed})
             # return the new task, to be started if not run before
             return pushed

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/types.py

@@ -11,6 +11,7 @@ from typing import (
     Literal,
     NamedTuple,
     Optional,
+    ParamSpec,
     TypeVar,
     Union,
     cast,
@@ -22,6 +23,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:
@@ -30,7 +32,7 @@ if TYPE_CHECKING:
 
 
 try:
-    from langchain_core.messages.tool import ToolOutputMixin
+    from langchain_core.messages.tool import ToolOutputMixin  # type: ignore
 except ImportError:
 
     class ToolOutputMixin:  # type: ignore[no-redef]
@@ -122,13 +124,16 @@ class RetryPolicy(NamedTuple):
     """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
 
 
-class CachePolicy(NamedTuple):
+P = ParamSpec("P")
+
+
+class CachePolicy(NamedTuple, Generic[P]):
     """Configuration for caching nodes.
 
     !!! version-added "Added in version 0.2.24."
     """
 
-    pass
+    key: Callable[P, str | bytes] = default_cache_key
 
 
 @dataclasses.dataclass(**_DC_KWARGS)

libs/langgraph/langgraph/utils/cache.py

@@ -0,0 +1,26 @@
+from collections.abc import Hashable
+from typing import Any
+
+
+def _freeze(obj: Any) -> Hashable:
+    if isinstance(obj, dict):
+        # sort keys so {"a":1,"b":2} == {"b":2,"a":1}
+        return tuple(sorted((k, _freeze(v)) for k, v in obj.items()))
+    elif isinstance(obj, (list, tuple, set, frozenset)):
+        return tuple(_freeze(x) 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) -> 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)