Add support for multiple subgraphs called in a single node

libs/langgraph/langgraph/errors.py

@@ -107,18 +107,3 @@ class CheckpointNotLatest(Exception):
     """Raised when the checkpoint is not the latest version (for distributed mode)."""
 
     pass
-
-
-class MultipleSubgraphsError(Exception):
-    """Raised when multiple subgraphs are called inside the same node.
-
-    Troubleshooting guides:
-
-    - [MULTIPLE_SUBGRAPHS](https://python.langchain.com/docs/troubleshooting/errors/MULTIPLE_SUBGRAPHS)
-    """
-
-    pass
-
-
-_SEEN_CHECKPOINT_NS: set[str] = set()
-"""Used for subgraph detection."""

libs/langgraph/langgraph/pregel/algo.py

@@ -766,6 +766,8 @@ def _scratchpad(
             (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME),
             MISSING,
         ),
+        # subgraph
+        subgraph_counter=0,
     )
 
 

libs/langgraph/langgraph/pregel/loop.py

@@ -63,12 +63,10 @@ from langgraph.constants import (
     TAG_HIDDEN,
 )
 from langgraph.errors import (
-    _SEEN_CHECKPOINT_NS,
     CheckpointNotLatest,
     EmptyInputError,
     GraphDelegate,
     GraphInterrupt,
-    MultipleSubgraphsError,
 )
 from langgraph.managed.base import (
     ManagedValueMapping,
@@ -116,6 +114,7 @@ from langgraph.types import (
     Command,
     LoopProtocol,
     PregelExecutableTask,
+    PregelScratchpad,
     StreamChunk,
     StreamProtocol,
 )
@@ -230,20 +229,26 @@ class PregelLoop(LoopProtocol):
         self.debug = debug
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
             self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
+        scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
+        if scratchpad is not None:
+            if scratchpad["subgraph_counter"]:
+                self.config = patch_configurable(
+                    self.config,
+                    {
+                        CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join(
+                            (
+                                config[CONF][CONFIG_KEY_CHECKPOINT_NS],
+                                str(scratchpad["subgraph_counter"]),
+                            )
+                        )
+                    },
+                )
+            scratchpad["subgraph_counter"] += 1
         if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
             self.config = patch_configurable(
                 self.config,
                 {CONFIG_KEY_CHECKPOINT_NS: "", CONFIG_KEY_CHECKPOINT_ID: None},
             )
-        if check_subgraphs and self.is_nested and self.checkpointer is not None:
-            if self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] in _SEEN_CHECKPOINT_NS:
-                raise MultipleSubgraphsError(
-                    "Multiple subgraphs called inside the same node\n\n"
-                    "Troubleshooting URL: https://python.langchain.com/docs"
-                    "/troubleshooting/errors/MULTIPLE_SUBGRAPHS/"
-                )
-            else:
-                _SEEN_CHECKPOINT_NS.add(self.config[CONF][CONFIG_KEY_CHECKPOINT_NS])
         if (
             CONFIG_KEY_CHECKPOINT_MAP in self.config[CONF]
             and self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)

libs/langgraph/langgraph/pregel/retry.py

@@ -12,7 +12,7 @@ from langgraph.constants import (
     CONFIG_KEY_RESUMING,
     NS_SEP,
 )
-from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
+from langgraph.errors import GraphBubbleUp, ParentCommand
 from langgraph.types import Command, PregelExecutableTask, RetryPolicy
 from langgraph.utils.config import patch_configurable
 
@@ -96,13 +96,6 @@ def run_with_retry(
             )
             # signal subgraphs to resume (if available)
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
-        finally:
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
 
 
 async def arun_with_retry(
@@ -188,10 +181,3 @@ async def arun_with_retry(
             )
             # signal subgraphs to resume (if available)
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
-        finally:
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)

libs/langgraph/langgraph/types.py

@@ -346,6 +346,8 @@ class PregelScratchpad(TypedDict):
     interrupt_counter: int
     resume: list[Any]
     null_resume: Any
+    # subgraph
+    subgraph_counter: int
 
 
 def interrupt(value: Any) -> Any:

libs/langgraph/tests/test_pregel.py

@@ -51,7 +51,7 @@ from langgraph.checkpoint.base import (
 )
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
-from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError
+from langgraph.errors import InvalidUpdateError
 from langgraph.func import entrypoint, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
@@ -1745,9 +1745,8 @@ def test_invoke_join_then_call_other_pregel(
 
     # add checkpointer
     app.checkpointer = checkpointer
-    # subgraph is called twice in the same node, through .map(), so raises
-    with pytest.raises(MultipleSubgraphsError):
-        app.invoke([2, 3], {"configurable": {"thread_id": "1"}})
+    # subgraph is called twice in the same node, but that works
+    assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
 
     # set inner graph checkpointer NeverCheckpoint
     inner_app.checkpointer = False

libs/langgraph/tests/test_pregel_async.py

@@ -48,7 +48,7 @@ from langgraph.checkpoint.base import (
 )
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
-from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
+from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.func import entrypoint, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessagesState, add_messages
@@ -4068,9 +4068,8 @@ async def test_invoke_join_then_call_other_pregel(
     async with awith_checkpointer(checkpointer_name) as checkpointer:
         # add checkpointer
         app.checkpointer = checkpointer
-        # subgraph is called twice in the same node, through .map(), so raises
-        with pytest.raises(MultipleSubgraphsError):
-            await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}})
+        # subgraph is called twice, and that works
+        assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
 
         # set inner graph checkpointer NeverCheckpoint
         inner_app.checkpointer = False