diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 45138de0f..f6a29198b 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6aabc9a61..f59fdd951 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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(): diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index be5c0a91c..42b73de36 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/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 diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 916abef36..b0b2f31a9 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/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, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 9ede99db1..53233a0ce 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 38c3ed52d..03d3f329a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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: