Fix

langchain-ai · Oct 15, 2024 · d48faec · d48faec
1 parent 5e175e0
commit d48faec
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
@@ -437,8 +437,8 @@ def _prepare_state_snapshot(
             saved.checkpoint,
             LoopProtocol(
                 config=saved.config,
-                step=saved.metadata["step"],
-                stop=saved.metadata["step"] + 1,
+                step=saved.metadata.get("step", -1) + 1,
+                stop=saved.metadata.get("step", -1) + 2,
             ),
             skip_context=True,
         ) as (channels, managed):
@@ -522,8 +522,8 @@ async def _aprepare_state_snapshot(
             saved.checkpoint,
             LoopProtocol(
                 config=saved.config,
-                step=saved.metadata["step"],
-                stop=saved.metadata["step"] + 1,
+                step=saved.metadata.get("step", -1) + 1,
+                stop=saved.metadata.get("step", -1) + 2,
             ),
             skip_context=True,
         ) as (
@@ -852,7 +852,7 @@ def update_state(
         with ChannelsManager(
             self.channels,
             checkpoint,
-            LoopProtocol(config=config, step=step, stop=step + 1),
+            LoopProtocol(config=config, step=step + 1, stop=step + 2),
         ) as (
             channels,
             managed,
@@ -1002,7 +1002,7 @@ async def aupdate_state(
         async with AsyncChannelsManager(
             self.channels,
             checkpoint,
-            LoopProtocol(config=config, step=step, stop=step + 1),
+            LoopProtocol(config=config, step=step + 1, stop=step + 2),
         ) as (
             channels,
             managed,

diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
@@ -191,7 +191,6 @@ async def attempt(self, msg: MessageToExecutor) -> None:
                 step=saved.metadata["step"] + 1,
                 stop=saved.metadata["step"] + 2,
             ),
-            self.graph.store,
         ) as (channels, managed), AsyncBackgroundExecutor() as submit:
             if task := await asyncio.to_thread(
                 prepare_single_task,