From 2742b2f8846e20ac3160baf5066abcc296b2dc8a Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 2 Aug 2024 12:59:15 -0700
Subject: [PATCH] aupdate_state now accepts null values (#1181)

* aupdate_state now accepts null values

---------

Co-authored-by: vbarda <vadym@langchain.dev>
---
 libs/langgraph/langgraph/pregel/__init__.py | 51 ++++++++++++---------
 libs/langgraph/pyproject.toml               |  2 +-
 libs/langgraph/tests/test_pregel_async.py   | 33 +++++++++++--
 3 files changed, 61 insertions(+), 25 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 3a247c579..6e81f2e90 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -627,8 +627,35 @@ class Pregel(
         # get last checkpoint
         saved = await self.checkpointer.aget_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
+        step = saved.metadata.get("step", -1) if saved else -1
+        # merge configurable fields with previous checkpoint config
+        checkpoint_config = {
+            **config,
+            "configurable": {
+                **config["configurable"],
+                # TODO: add proper support for updating nested subgraph state
+                "checkpoint_ns": "",
+            },
+        }
+        if saved:
+            checkpoint_config = {
+                "configurable": {
+                    **config.get("configurable", {}),
+                    **saved.config["configurable"],
+                }
+            }
         # find last node that updated the state, if not provided
-        if as_node is None and not saved:
+        if values is None and as_node is None:
+            return await self.checkpointer.aput(
+                checkpoint_config,
+                create_checkpoint(checkpoint, None, step),
+                {
+                    "source": "update",
+                    "step": step,
+                    "writes": {},
+                },
+            )
+        elif as_node is None and not saved:
             if (
                 isinstance(self.input_channels, str)
                 and self.input_channels in self.nodes
@@ -685,31 +712,13 @@ class Pregel(
             apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             )
-            step = saved.metadata.get("step", -2) + 1 if saved else -1
-
-            # merge configurable fields with previous checkpoint config
-            checkpoint_config = {
-                **config,
-                "configurable": {
-                    **config["configurable"],
-                    # TODO: add proper support for updating nested subgraph state
-                    "checkpoint_ns": "",
-                },
-            }
-            if saved:
-                checkpoint_config = {
-                    "configurable": {
-                        **config.get("configurable", {}),
-                        **saved.config["configurable"],
-                    }
-                }
 
             return await self.checkpointer.aput(
                 checkpoint_config,
-                create_checkpoint(checkpoint, channels, step),
+                create_checkpoint(checkpoint, channels, step + 1),
                 {
                     "source": "update",
-                    "step": step,
+                    "step": step + 1,
                     "writes": {as_node: values},
                 },
             )
diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml
index 828bdb435..f02e6123d 100644
--- a/libs/langgraph/pyproject.toml
+++ b/libs/langgraph/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.1.17"
+version = "0.1.19"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index b718f51c9..c646e21a2 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -853,6 +853,33 @@ async def test_invoke_two_processes_in_out_interrupt(
                 parent_config=None,
             ),
         ]
+
+        # forking from any previous checkpoint w/out forking should do nothing
+        assert [
+            c async for c in app.astream(None, history[0].config, stream_mode="updates")
+        ] == []
+        assert [
+            c async for c in app.astream(None, history[1].config, stream_mode="updates")
+        ] == []
+        assert [
+            c async for c in app.astream(None, history[2].config, stream_mode="updates")
+        ] == []
+
+        # forking and re-running from any prev checkpoint should re-run nodes
+        fork_config = await app.aupdate_state(history[0].config, None)
+        assert [
+            c async for c in app.astream(None, fork_config, stream_mode="updates")
+        ] == []
+
+        fork_config = await app.aupdate_state(history[1].config, None)
+        assert [
+            c async for c in app.astream(None, fork_config, stream_mode="updates")
+        ] == [{"two": {"output": 5}}]
+
+        fork_config = await app.aupdate_state(history[2].config, None)
+        assert [
+            c async for c in app.astream(None, fork_config, stream_mode="updates")
+        ] == [{"one": {"inbox": 4}}]
     finally:
         if hasattr(checkpointer, "__aexit__"):
             await checkpointer.__aexit__(None, None, None)
@@ -5252,7 +5279,7 @@ async def test_branch_then() -> None:
             created_at=AnyStr(),
             metadata={
                 "source": "update",
-                "step": -1,
+                "step": 0,
                 "writes": {START: {"my_key": "key", "market": "DE"}},
             },
         )
@@ -5271,7 +5298,7 @@ async def test_branch_then() -> None:
             ],
             metadata={
                 "source": "loop",
-                "step": 0,
+                "step": 1,
                 "writes": {"prepare": {"my_key": " prepared"}},
             },
             parent_config=uconfig,
@@ -5290,7 +5317,7 @@ async def test_branch_then() -> None:
             ],
             metadata={
                 "source": "loop",
-                "step": 2,
+                "step": 3,
                 "writes": {"finish": {"my_key": " finished"}},
             },
             parent_config=[