Merge pull request #2627 from langchain-ai/nc/4dec/state-ensure-config

lib: Call ensure_config in state crud methods
This commit is contained in:
Nuno Campos
2024-12-04 11:53:03 -05:00
committed by GitHub
+8 -6
View File
@@ -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: