diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 62506142d..e714afe21 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -673,7 +673,7 @@ class Pregel(PregelProtocol): self, config: RunnableConfig, *, subgraphs: bool = False ) -> StateSnapshot: """Get the current state of the graph.""" - checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get( + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: @@ -710,7 +710,7 @@ class Pregel(PregelProtocol): self, config: RunnableConfig, *, subgraphs: bool = False ) -> StateSnapshot: """Get the current state of the graph.""" - checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get( + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: @@ -751,8 +751,9 @@ class Pregel(PregelProtocol): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[StateSnapshot]: + config = ensure_config(config) """Get the history of the state of the graph.""" - checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get( + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: @@ -800,8 +801,9 @@ class Pregel(PregelProtocol): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> AsyncIterator[StateSnapshot]: + config = ensure_config(config) """Get the history of the state of the graph.""" - checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get( + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: @@ -855,7 +857,7 @@ class Pregel(PregelProtocol): node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. """ - checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get( + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: @@ -1130,7 +1132,7 @@ class Pregel(PregelProtocol): values: dict[str, Any] | Any, as_node: Optional[str] = None, ) -> RunnableConfig: - checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get( + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: