diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 143052d0f..d83eb906b 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -799,9 +799,15 @@ class Pregel( and self.checkpointer.at == CheckpointAt.END_OF_RUN ): checkpoint = create_checkpoint(checkpoint, channels) - checkpoint_config = self.checkpointer.put( - checkpoint_config, checkpoint + executor.submit( + self.checkpointer.put(checkpoint_config, checkpoint) ) + checkpoint_config = { + "configurable": { + "thread_id": checkpoint_config["configurable"]["thread_id"], + "thread_ts": checkpoint["ts"], + } + } if stream_mode == "debug": yield map_debug_checkpoint( step, @@ -854,6 +860,7 @@ class Pregel( None, ) try: + tasks: list[asyncio.Task] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -1058,9 +1065,17 @@ class Pregel( and self.checkpointer.at == CheckpointAt.END_OF_RUN ): checkpoint = create_checkpoint(checkpoint, channels) - checkpoint_config = await self.checkpointer.aput( - checkpoint_config, checkpoint + tasks.append( + asyncio.create_task( + self.checkpointer.aput(checkpoint_config, checkpoint) + ) ) + checkpoint_config = { + "configurable": { + "thread_id": checkpoint_config["configurable"]["thread_id"], + "thread_ts": checkpoint["ts"], + } + } if stream_mode == "debug": yield map_debug_checkpoint( step, checkpoint_config, channels, self.stream_channels_asis @@ -1077,8 +1092,11 @@ class Pregel( try: for task in futures: task.cancel() + tasks.append(task) except NameError: pass + # wait for all tasks to finish + await asyncio.gather(*tasks, return_exceptions=True) def invoke( self,