From 0822a287e37445fc81b121dcef468534ce946118 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 11 Oct 2024 11:24:56 -0700 Subject: [PATCH 1/2] fix: Work w python's O flag - assert statements are skipped in that case, so we need to move calls to apply_writes to outside assert statements --- libs/langgraph/langgraph/pregel/__init__.py | 10 ++++++---- libs/langgraph/langgraph/pregel/loop.py | 5 +++-- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 98d82712b..0b81e5b77 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -866,9 +866,10 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): if saved: checkpointer.put_writes(checkpoint_config, task.writes, task_id) # apply to checkpoint and save - assert not apply_writes( + mv_writes = apply_writes( checkpoint, channels, [task], checkpointer.get_next_version - ), "Can't write to SharedValues from update_state" + ) + assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) next_config = checkpointer.put( checkpoint_config, @@ -1011,9 +1012,10 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): if saved: await checkpointer.aput_writes(checkpoint_config, writes, task_id) # apply to checkpoint and save - assert not apply_writes( + mv_writes = not apply_writes( checkpoint, channels, [task], checkpointer.get_next_version - ), "Can't write to SharedValues from update_state" + ) + assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) next_config = await checkpointer.aput( checkpoint_config, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 2ac371f2a..b3bffa544 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -503,12 +503,13 @@ class PregelLoop: manager=None, ) # apply input writes - assert not apply_writes( + mv_writes = apply_writes( self.checkpoint, self.channels, [*discard_tasks.values(), PregelTaskWrites(INPUT, input_writes, [])], self.checkpointer_get_next_version, - ), "Can't write to SharedValues in graph input" + ) + assert not mv_writes, "Can't write to SharedValues in graph input" # save input checkpoint self._put_checkpoint({"source": "input", "writes": dict(input_writes)}) elif CONFIG_KEY_RESUMING not in configurable: From fc20de5bbaa2b9e5b35a7201dace46e88477f971 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 11 Oct 2024 11:39:19 -0700 Subject: [PATCH 2/2] Fix --- libs/langgraph/langgraph/pregel/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 0b81e5b77..441647355 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1012,7 +1012,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): if saved: await checkpointer.aput_writes(checkpoint_config, writes, task_id) # apply to checkpoint and save - mv_writes = not apply_writes( + mv_writes = apply_writes( checkpoint, channels, [task], checkpointer.get_next_version ) assert not mv_writes, "Can't write to SharedValues from update_state"