From 904a1a34716a9b4c4e78468c5266b1d135b8bd72 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 18:09:10 -0400 Subject: [PATCH] extra paranoia --- libs/langgraph/langgraph/pregel/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2173cc7d5..496ec1b30 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -223,11 +223,12 @@ def _prepare_state_snapshot( checkpoint_tuples: Iterator[CheckpointTuple], ) -> StateSnapshot: checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id") checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} for saved in checkpoint_tuples: saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] - if saved_checkpoint_id != config["configurable"]["checkpoint_id"]: + if checkpoint_id and saved_checkpoint_id != checkpoint_id: continue graph_checkpoint_ns = saved_checkpoint_ns.split( @@ -287,11 +288,12 @@ async def _prepare_state_snapshot_async( checkpoint_tuples: AsyncIterator[CheckpointTuple], ) -> StateSnapshot: checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id") checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} async for saved in checkpoint_tuples: saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] - if saved_checkpoint_id != config["configurable"]["checkpoint_id"]: + if checkpoint_id and saved_checkpoint_id != checkpoint_id: continue graph_checkpoint_ns = saved_checkpoint_ns.split(