From 6e0041529e44e547fa954a486cd409ff703acd1d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 16:25:59 -0700 Subject: [PATCH] Add clear cache methods --- libs/langgraph/langgraph/func/__init__.py | 63 ++++++++++++++------- libs/langgraph/langgraph/pregel/__init__.py | 40 +++++++++++++ libs/langgraph/tests/test_pregel.py | 53 ++++++++++++++--- libs/langgraph/tests/test_pregel_async.py | 43 ++++++++++++++ 4 files changed, 170 insertions(+), 29 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 06aac25c2..de0e25648 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/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__) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f59fdd951..286a98ddc 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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.""" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 53233a0ce..bca22ceb2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c43f4cdac..63c0a4b3e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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)