From d6492ef048618bbd9713aba2968041fcfa1140ed Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 15:32:48 -0800 Subject: [PATCH] Add support for multiple subgraphs called in a single node --- libs/langgraph/langgraph/errors.py | 15 ------------- libs/langgraph/langgraph/pregel/algo.py | 2 ++ libs/langgraph/langgraph/pregel/loop.py | 27 ++++++++++++++--------- libs/langgraph/langgraph/pregel/retry.py | 16 +------------- libs/langgraph/langgraph/types.py | 2 ++ libs/langgraph/tests/test_pregel.py | 7 +++--- libs/langgraph/tests/test_pregel_async.py | 7 +++--- 7 files changed, 27 insertions(+), 49 deletions(-) diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 0737a31d0..8e78a8784 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/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.""" diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c46138550..205793ab4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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, ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 9a745a2e7..ff244b7de 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 29faaab21..43e7e8d9e 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/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) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3b1bcd213..09f777d01 100644 --- a/libs/langgraph/langgraph/types.py +++ b/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: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 67536deed..76a067533 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 2f7d7a1c4..a6fab066f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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