Finish implementation, add tests

libs/langgraph/langgraph/graph/state.py

@@ -80,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
@@ -114,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
 
@@ -260,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.
@@ -277,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."""
@@ -291,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.
@@ -304,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.
@@ -432,6 +437,7 @@ class StateGraph(Graph):
             metadata,
             input=input or self.schema,
             retry_policy=retry,
+            cache_policy=cache_policy,
             ends=ends,
             defer=defer,
         )
@@ -814,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

@@ -2827,7 +2827,7 @@ class Pregel(PregelProtocol):
                         [t for t in loop.tasks.values() if not t.writes],
                         timeout=self.step_timeout,
                         get_waiter=get_waiter,
-                        # TODO pass match_cached_writes
+                        match_cached_writes=loop.amatch_cached_writes,
                     ):
                         # emit output
                         for o in output():

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

libs/langgraph/langgraph/pregel/runner.py

@@ -155,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(
@@ -275,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)
@@ -286,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(
@@ -301,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,
@@ -348,6 +356,7 @@ class PregelRunner:
                             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,
@@ -609,7 +618,7 @@ def _acall(
     input: Any,
     *,
     retry: Optional[Sequence[RetryPolicy]] = None,
-    cache: Optional[CachePolicy] = None,
+    cache_policy: Optional[CachePolicy] = None,
     callbacks: Callbacks = None,
     # injected dependencies
     futures: weakref.ref[FuturesDict],
@@ -618,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,
@@ -630,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, cache_policy=cache, callbacks=callbacks),
+        Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks),
     ):
         if fut := next(
             (
@@ -666,6 +678,7 @@ def _acall(
                     next_task,
                     retry,
                     stream=stream,
+                    match_cached_writes=match_cached_writes,
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _acall,
@@ -673,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/tests/test_pregel.py

@@ -3573,7 +3573,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, file_cache: BaseCache
+) -> None:
     def sorted_add(
         x: list[str], y: Union[list[str], list[tuple[str, str]]]
     ) -> list[str]:
@@ -3588,7 +3591,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:
@@ -3615,7 +3622,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)
@@ -3630,7 +3641,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=file_cache)
 
     assert app.invoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: analyzed: query: what is weather in sf",
@@ -3639,12 +3650,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"
@@ -3655,6 +3678,7 @@ 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
 
 
 def test_callable_in_conditional_edges_with_no_path_map() -> None:

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, file_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=file_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,7 @@ 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
 
 
 async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
@@ -7510,6 +7537,47 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None:
         assert counter == 3
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_multiple_interrupts_functional_cache(
+    checkpointer_name: str, file_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=file_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
+
+
 @NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: