Merge pull request #2087 from langchain-ai/nc/11oct/o-flag

fix: Work w python's O flag
langchain-ai · Oct 11, 2024 · c9adf99 · c9adf99
2 parents 66741ba + fc20de5
commit c9adf99
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
@@ -866,9 +866,10 @@ def update_state(
             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 @@ async def aupdate_state(
             if saved:
                 await checkpointer.aput_writes(checkpoint_config, 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 = await checkpointer.aput(
                 checkpoint_config,

diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
@@ -503,12 +503,13 @@ def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
                 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: