diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 2c46108c3..88c31f74c 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -439,7 +439,9 @@ class PregelLoop: return None return self._graph_lifecycle_events.popleft() - def put_writes(self, task_id: str, writes: WritesT) -> None: + def put_writes( + self, task_id: str, writes: WritesT, *, cached: bool = False + ) -> None: """Put writes for a task, to be read by the next tick.""" if not writes: return @@ -532,7 +534,7 @@ class PregelLoop: self._error_handler_write_futs.append(fut) # output writes if hasattr(self, "tasks"): - self.output_writes(task_id, writes) + self.output_writes(task_id, writes, cached=cached) def _put_pending_writes(self) -> None: if self.checkpointer_put_writes is None: @@ -1684,7 +1686,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) -> PregelExecutableTask | None: if pushed := super().accept_push(task, write_idx, call): for task in self.match_cached_writes(): - self.output_writes(task.id, task.writes, cached=True) + self.put_writes(task.id, task.writes, cached=True) return pushed def schedule_error_handler( @@ -1721,13 +1723,15 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): if self._reapplies_pending_writes: self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task}) for task in self.match_cached_writes(): - self.output_writes(task.id, task.writes, cached=True) + self.put_writes(task.id, task.writes, cached=True) return handler_task - def put_writes(self, task_id: str, writes: WritesT) -> None: + def put_writes( + self, task_id: str, writes: WritesT, *, cached: bool = False + ) -> None: """Put writes for a task, to be read by the next tick.""" - super().put_writes(task_id, writes) - if not writes or self.cache is None or not hasattr(self, "tasks"): + super().put_writes(task_id, writes, cached=cached) + if cached or not writes or self.cache is None or not hasattr(self, "tasks"): return task = self.tasks.get(task_id) if task is None or task.cache_key is None: @@ -1937,7 +1941,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): ) -> PregelExecutableTask | None: if pushed := super().accept_push(task, write_idx, call): for task in await self.amatch_cached_writes(): - self.output_writes(task.id, task.writes, cached=True) + self.put_writes(task.id, task.writes, cached=True) return pushed async def aschedule_error_handler( @@ -1974,13 +1978,15 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if self._reapplies_pending_writes: self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task}) for task in await self.amatch_cached_writes(): - self.output_writes(task.id, task.writes, cached=True) + self.put_writes(task.id, task.writes, cached=True) return handler_task - def put_writes(self, task_id: str, writes: WritesT) -> None: + def put_writes( + self, task_id: str, writes: WritesT, *, cached: bool = False + ) -> None: """Put writes for a task, to be read by the next tick.""" - super().put_writes(task_id, writes) - if not writes or self.cache is None or not hasattr(self, "tasks"): + super().put_writes(task_id, writes, cached=cached) + if cached or not writes or self.cache is None or not hasattr(self, "tasks"): return task = self.tasks.get(task_id) if task is None or task.cache_key is None: diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1cf1aa4bb..48935dae0 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -3038,7 +3038,7 @@ class Pregel( # with channel updates applied only at the transition between steps. while loop.tick(): for task in loop.match_cached_writes(): - loop.output_writes(task.id, task.writes, cached=True) + loop.put_writes(task.id, task.writes, cached=True) for _ in runner.tick( [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, @@ -3509,7 +3509,7 @@ class Pregel( try: while loop.tick(): for task in await loop.amatch_cached_writes(): - loop.output_writes(task.id, task.writes, cached=True) + loop.put_writes(task.id, task.writes, cached=True) async for _ in runner.atick( [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, diff --git a/libs/langgraph/tests/test_delta_channel_node_cache.py b/libs/langgraph/tests/test_delta_channel_node_cache.py new file mode 100644 index 000000000..d0171f807 --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_node_cache.py @@ -0,0 +1,89 @@ +"""A node served from the node cache saves its writes like a node that ran, +without writing them back to the cache.""" + +import operator +from typing import Annotated, Any + +import pytest +from langgraph.cache.memory import InMemoryCache +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import TypedDict + +from langgraph.channels.delta import DeltaChannel +from langgraph.graph import START, StateGraph +from langgraph.types import CachePolicy, Durability + +pytestmark = pytest.mark.anyio + +INPUT = {"log": [], "plain": []} + + +def _append(current: list, writes: list) -> list: + return [*current, *(item for write in writes for item in write)] + + +class _State(TypedDict): + log: Annotated[list, DeltaChannel(_append)] + plain: Annotated[list, operator.add] + + +class _CountsSets(InMemoryCache): + sets = 0 + + def set(self, keys: Any) -> None: + self.sets += 1 + super().set(keys) + + +def _a_then_cached_b_then_c(runs: list[str], cache: InMemoryCache) -> Any: + def node(name: str) -> Any: + def run(state: _State) -> dict: + runs.append(name) + return {"log": [name], "plain": [name]} + + return run + + builder = StateGraph(_State) + builder.add_node("a", node("a")) + builder.add_node("b", node("b"), cache_policy=CachePolicy()) + builder.add_node("c", node("c")) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("b", "c") + return builder.compile(checkpointer=InMemorySaver(), cache=cache) + + +def test_a_cache_hit_saves_its_writes_without_caching_them_again( + durability: Durability, +) -> None: + runs: list[str] = [] + cache = _CountsSets() + graph = _a_then_cached_b_then_c(runs, cache) + graph.invoke(INPUT, {"configurable": {"thread_id": "1"}}, durability=durability) + config = {"configurable": {"thread_id": "2"}} + + graph.invoke(INPUT, config, durability=durability) + + assert runs == ["a", "b", "c", "a", "c"] + assert cache.sets == 1, "the cache hit was written back to the cache" + for state in graph.get_state_history(config): + assert state.values.get("log", []) == state.values.get("plain", []) + + +async def test_a_cache_hit_saves_its_writes_without_caching_them_again_async( + durability: Durability, +) -> None: + runs: list[str] = [] + cache = _CountsSets() + graph = _a_then_cached_b_then_c(runs, cache) + await graph.ainvoke( + INPUT, {"configurable": {"thread_id": "1"}}, durability=durability + ) + config = {"configurable": {"thread_id": "2"}} + + await graph.ainvoke(INPUT, config, durability=durability) + + assert runs == ["a", "b", "c", "a", "c"] + assert cache.sets == 1, "the cache hit was written back to the cache" + async for state in graph.aget_state_history(config): + assert state.values.get("log", []) == state.values.get("plain", []) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bd76c5ea7..d698da84f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5932,9 +5932,9 @@ def test_no_redundant_put_writes_for_cached_task( put_writes_task_ids: list[str] = [] orig = PregelLoop.put_writes - def spy(self, task_id, writes): + def spy(self, task_id, writes, **kwargs): put_writes_task_ids.append(task_id) - return orig(self, task_id, writes) + return orig(self, task_id, writes, **kwargs) with patch.object(PregelLoop, "put_writes", spy): result = workflow.invoke(Command(resume="ans"), config=config) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 150801ffc..4c40356d7 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8157,9 +8157,9 @@ async def test_no_redundant_put_writes_for_cached_task( put_writes_task_ids: list[str] = [] orig = PregelLoop.put_writes - def spy(self, task_id, writes): + def spy(self, task_id, writes, **kwargs): put_writes_task_ids.append(task_id) - return orig(self, task_id, writes) + return orig(self, task_id, writes, **kwargs) with patch.object(PregelLoop, "put_writes", spy): result = await workflow.ainvoke(Command(resume="ans"), config=config)