Add clear cache methods

libs/langgraph/langgraph/func/__init__.py

@@ -20,7 +20,7 @@ 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,
@@ -28,6 +28,7 @@ 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
@@ -35,6 +36,40 @@ from langgraph.store.base import BaseStore
 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:
+        self.func = func
+        self.retry = retry
+        self.cache_policy = cache_policy
+        functools.update_wrapper(self, func)
+        if name is not None:
+            setattr(self, "__name__", name)
+
+    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.delete(((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.adelete(
+                ((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),)
+            )
+
+
 @overload
 def task(
     *,
@@ -43,14 +78,14 @@ def task(
     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(
@@ -62,9 +97,9 @@ def task(
 ) -> 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.
 
@@ -132,23 +167,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, cache_policy=cache_policy
+        return TaskFunction(
+            func, retry=retry_policies, cache_policy=cache_policy, name=name
         )
-        object.__setattr__(call_func, "_is_pregel_task", True)
-        return functools.update_wrapper(call_func, func)
 
     if __func_or_none__ is not None:
         return decorator(__func_or_none__)

libs/langgraph/langgraph/pregel/__init__.py

@@ -47,6 +47,7 @@ from langgraph.checkpoint.base import (
     copy_checkpoint,
 )
 from langgraph.constants import (
+    CACHE_NS_WRITES,
     CONF,
     CONFIG_KEY_CACHE,
     CONFIG_KEY_CHECKPOINT_DURING,
@@ -87,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
@@ -2989,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.delete(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.adelete(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/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
@@ -3680,6 +3673,17 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple(
     ]
     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:
     class State(TypedDict, total=False):
@@ -6663,6 +6667,39 @@ def test_multiple_interrupts_functional_cache(
     }
     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(file_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(

libs/langgraph/tests/test_pregel_async.py

@@ -5505,6 +5505,17 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple(
     ]
     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:
     def sorted_add(
@@ -7578,6 +7589,38 @@ async def test_multiple_interrupts_functional_cache(
         }
         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(file_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)