From 1f29925034d788b3b8b8b8f26bcd0eadee3da286 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 14:44:42 -0400 Subject: [PATCH] add max recursion depth --- libs/langgraph/langgraph/pregel/__init__.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f6ce5aefb..33bae3330 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -182,10 +182,16 @@ def _get_checkpoint_ns_to_graph( graph: Pregel, checkpoint_ns_to_graph: Optional[dict[str, Pregel]] = None, checkpoint_ns: str = "", + max_depth: int = 10, ) -> Pregel: if checkpoint_ns_to_graph is None: checkpoint_ns_to_graph = {} + if max_depth <= 0: + raise RecursionError( + f"Reached maximum recursion depth while building checkpoint NS -> graph mapping." + ) + for node_name, node in graph.nodes.items(): new_checkpoint_ns = ( f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" @@ -194,13 +200,16 @@ def _get_checkpoint_ns_to_graph( ) if isinstance(node.bound, Pregel): _get_checkpoint_ns_to_graph( - node.bound, checkpoint_ns_to_graph, new_checkpoint_ns + node.bound, checkpoint_ns_to_graph, new_checkpoint_ns, max_depth - 1 ) 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, new_checkpoint_ns + runnable, + checkpoint_ns_to_graph, + new_checkpoint_ns, + max_depth - 1, ) checkpoint_ns_to_graph[checkpoint_ns] = graph