From 32952747116d769ecbdb51dc2d39fc7ceef91c55 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 13 Aug 2024 10:23:53 -0400 Subject: [PATCH] fix empty snapshot --- libs/langgraph/langgraph/pregel/__init__.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 0b7dd76c9..58e34dd59 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -499,7 +499,8 @@ class Pregel( if include_subgraph_state: checkpoint_tuples = self.checkpointer.list(config) else: - checkpoint_tuples = iter([self.checkpointer.get_tuple(config)]) + checkpoint_tuple = self.checkpointer.get_tuple(config) + checkpoint_tuples = iter([checkpoint_tuple] if checkpoint_tuple else []) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = config["configurable"].get("checkpoint_id") @@ -544,7 +545,7 @@ class Pregel( if not checkpoint_ns_to_state_snapshots: return StateSnapshot( - values={}, next=(), config=config, checkpoint=empty_checkpoint() + values={}, next=(), config=config, metadata=None, created_at=None ) state_snapshot = self._assemble_state_snapshot_hierarchy( @@ -564,7 +565,9 @@ class Pregel( else: async def alist_checkpoints(): - yield await self.checkpointer.aget_tuple(config) + checkpoint_tuple = await self.checkpointer.aget_tuple(config) + if checkpoint_tuple: + yield checkpoint_tuple checkpoint_tuples = alist_checkpoints() @@ -611,7 +614,7 @@ class Pregel( if not checkpoint_ns_to_state_snapshots: return StateSnapshot( - values={}, next=(), config=config, checkpoint=empty_checkpoint() + values={}, next=(), config=config, metadata=None, created_at=None ) state_snapshot = self._assemble_state_snapshot_hierarchy(