From 1333d8b4781ecad144b220c6427dbe1f8fb27d50 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 14:26:49 -0400 Subject: [PATCH] cleanup --- libs/langgraph/langgraph/pregel/__init__.py | 24 +++++++++++---------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2ee00bb29..f6ce5aefb 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -179,26 +179,28 @@ class Channel: def _get_checkpoint_ns_to_graph( - graph: Pregel, checkpoint_ns_to_graph: dict[str, Pregel] = {}, checkpoint_ns="" + graph: Pregel, + checkpoint_ns_to_graph: Optional[dict[str, Pregel]] = None, + checkpoint_ns: str = "", ) -> Pregel: + if checkpoint_ns_to_graph is None: + checkpoint_ns_to_graph = {} + for node_name, node in graph.nodes.items(): + new_checkpoint_ns = ( + f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" + if checkpoint_ns + else node_name + ) if isinstance(node.bound, Pregel): _get_checkpoint_ns_to_graph( - node.bound, - checkpoint_ns_to_graph, - f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" - if checkpoint_ns - else node_name, + node.bound, checkpoint_ns_to_graph, new_checkpoint_ns ) elif isinstance(node.bound, RunnableSequence): for runnable in node.bound.steps: if isinstance(runnable, Pregel): _get_checkpoint_ns_to_graph( - runnable, - checkpoint_ns_to_graph, - f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" - if checkpoint_ns - else node_name, + runnable, checkpoint_ns_to_graph, new_checkpoint_ns ) checkpoint_ns_to_graph[checkpoint_ns] = graph