From f9409022ed968293f4354c2726a8e00653721886 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 4 Nov 2024 15:10:12 -0800 Subject: [PATCH 1/3] lib: In calls to get_state apply pending writes - When calling get_state without a checkpoint id (ie to get the latest state) apply any pending writes for current checkpoint --- libs/langgraph/langgraph/pregel/__init__.py | 35 ++++++++++++++++++--- libs/langgraph/tests/test_pregel.py | 7 ++++- libs/langgraph/tests/test_pregel_async.py | 10 ++++-- 3 files changed, 45 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2525081d1..d146a9529 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -54,6 +54,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.constants import ( CONF, + CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_NODE_FINISHED, @@ -439,6 +440,7 @@ class Pregel(PregelProtocol): config: RunnableConfig, saved: Optional[CheckpointTuple], recurse: Optional[BaseCheckpointSaver] = None, + apply_pending_writes: bool = False, ) -> StateSnapshot: if not saved: return StateSnapshot( @@ -469,7 +471,10 @@ class Pregel(PregelProtocol): managed, saved.config, saved.metadata.get("step", -1) + 1, - for_execution=False, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer or None, + manager=None, ) # get the subgraphs subgraphs = dict(self.get_subgraphs()) @@ -503,6 +508,12 @@ class Pregel(PregelProtocol): task_states[task.id] = subgraphs[task.name].get_state( config, subgraphs=True ) + # apply pending writes + if apply_pending_writes and saved.pending_writes: + for tid, *t in saved.pending_writes: + next_tasks[tid].writes.append(t) + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes(saved.checkpoint, channels, tasks, None) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), @@ -524,6 +535,7 @@ class Pregel(PregelProtocol): config: RunnableConfig, saved: Optional[CheckpointTuple], recurse: Optional[BaseCheckpointSaver] = None, + apply_pending_writes: bool = False, ) -> StateSnapshot: if not saved: return StateSnapshot( @@ -557,7 +569,10 @@ class Pregel(PregelProtocol): managed, saved.config, saved.metadata.get("step", -1) + 1, - for_execution=False, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer or None, + manager=None, ) # get the subgraphs subgraphs = {n: g async for n, g in self.aget_subgraphs()} @@ -591,6 +606,12 @@ class Pregel(PregelProtocol): task_states[task.id] = await subgraphs[task.name].aget_state( config, subgraphs=True ) + # apply pending writes + if apply_pending_writes and saved.pending_writes: + for tid, *t in saved.pending_writes: + next_tasks[tid].writes.append(t) + if tasks := [t for t in next_tasks.values() if t.writes]: + apply_writes(saved.checkpoint, channels, tasks, None) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), @@ -638,7 +659,10 @@ class Pregel(PregelProtocol): config = merge_configs(self.config, config) if self.config else config saved = checkpointer.get_tuple(config) return self._prepare_state_snapshot( - config, saved, recurse=checkpointer if subgraphs else None + config, + saved, + recurse=checkpointer if subgraphs else None, + apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF], ) async def aget_state( @@ -672,7 +696,10 @@ class Pregel(PregelProtocol): config = merge_configs(self.config, config) if self.config else config saved = await checkpointer.aget_tuple(config) return await self._aprepare_state_snapshot( - config, saved, recurse=checkpointer if subgraphs else None + config, + saved, + recurse=checkpointer if subgraphs else None, + apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF], ) def get_state_history( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8536a08c3..a8a614408 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1513,9 +1513,10 @@ def test_pending_writes_resume( assert two.calls == 2 # two attempts # latest checkpoint should be before nodes "one", "two" + # but we should have applied the write from "one" state = graph.get_state(thread1) assert state is not None - assert state.values == {"value": 1} + assert state.values == {"value": 3} assert state.next == ("one", "two") assert state.tasks == ( PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), @@ -1528,6 +1529,10 @@ def test_pending_writes_resume( "writes": None, "thread_id": "1", } + # get_state with checkpoint_id should not apply any pending writes + state = graph.get_state(state.config) + assert state is not None + assert state.values == {"value": 1} # should contain pending write of "one" checkpoint = checkpointer.get_tuple(thread1) assert checkpoint is not None diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d8299c930..ec4179c71 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -482,10 +482,11 @@ async def test_cancel_graph_astream(checkpointer_name: str) -> None: assert awhile.started is False # checkpoint with output of "alittlewhile" should not be saved + # but we should have applied pending writes if checkpointer is not None: state = await graph.aget_state(thread1) assert state is not None - assert state.values == {"value": 1} + assert state.values == {"value": 3} # 1 + 2 assert state.next == ( "aparallelwhile", "alittlewhile", @@ -1722,9 +1723,10 @@ async def test_pending_writes_resume( assert two.calls == 2 # latest checkpoint should be before nodes "one", "two" + # but we should have applied pending writes from "one" state = await graph.aget_state(thread1) assert state is not None - assert state.values == {"value": 1} + assert state.values == {"value": 3} assert state.next == ("one", "two") assert state.tasks == ( PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), @@ -1742,6 +1744,10 @@ async def test_pending_writes_resume( "writes": None, "thread_id": "1", } + # get_state with checkpoint_id should not apply any pending writes + state = await graph.aget_state(state.config) + assert state is not None + assert state.values == {"value": 1} # should contain pending write of "one" checkpoint = await checkpointer.aget_tuple(thread1) assert checkpoint is not None From d28734f2875e7aee34f777fc16811c0ab9eddbc9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 4 Nov 2024 15:12:34 -0800 Subject: [PATCH 2/3] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d146a9529..9aa3bc915 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -511,7 +511,7 @@ class Pregel(PregelProtocol): # apply pending writes if apply_pending_writes and saved.pending_writes: for tid, *t in saved.pending_writes: - next_tasks[tid].writes.append(t) + next_tasks[tid].writes.append(t) # type: ignore[arg-type] if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes(saved.checkpoint, channels, tasks, None) # assemble the state snapshot @@ -609,7 +609,7 @@ class Pregel(PregelProtocol): # apply pending writes if apply_pending_writes and saved.pending_writes: for tid, *t in saved.pending_writes: - next_tasks[tid].writes.append(t) + next_tasks[tid].writes.append(t) # type: ignore[arg-type] if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes(saved.checkpoint, channels, tasks, None) # assemble the state snapshot From de3b654735104e460320afabf0d49f20c8d6d8f7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 4 Nov 2024 15:14:29 -0800 Subject: [PATCH 3/3] Add one more assertion --- libs/langgraph/tests/test_pregel_async.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index ec4179c71..b9d7936db 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2123,6 +2123,8 @@ async def test_max_concurrency(checkpointer_name: str) -> None: thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} assert await graph.ainvoke(["0"], thread1) == ["0", "1"] + state = await graph.aget_state(thread1) + assert state.values == ["0", "1"] assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"]