Guard and cache calls to find_subgraph_pregel

- The result of these doesnt change once a node is created, and it's fairly expensive to run, so great thing to cache
- There's a variety of errors that can come from inspecting the source code of a function (part of what this does) so adding a catch-all try-except block as this should be best-effort, not crash your graph

libs/langgraph/langgraph/pregel/__init__.py

@@ -97,7 +97,7 @@ from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.runner import PregelRunner
-from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_versions
+from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
@@ -429,7 +429,7 @@ class Pregel(PregelProtocol):
 
     def get_subgraphs(
         self, *, namespace: Optional[str] = None, recurse: bool = False
-    ) -> Iterator[tuple[str, Pregel]]:
+    ) -> Iterator[tuple[str, PregelProtocol]]:
         for name, node in self.nodes.items():
             # filter by prefix
             if namespace is not None:
@@ -437,7 +437,7 @@ class Pregel(PregelProtocol):
                     continue
 
             # find the subgraph, if any
-            graph = cast(Optional[Pregel], find_subgraph_pregel(node.bound))
+            graph = node.subgraphs[0] if node.subgraphs else None
 
             # if found, yield recursively
             if graph:
@@ -446,7 +446,7 @@ class Pregel(PregelProtocol):
                     return  # we found it, stop searching
                 if namespace is None:
                     yield name, graph
-                if recurse:
+                if recurse and isinstance(graph, Pregel):
                     if namespace is not None:
                         namespace = namespace[len(name) + 1 :]
                     yield from (
@@ -458,7 +458,7 @@ class Pregel(PregelProtocol):
 
     async def aget_subgraphs(
         self, *, namespace: Optional[str] = None, recurse: bool = False
-    ) -> AsyncIterator[tuple[str, Pregel]]:
+    ) -> AsyncIterator[tuple[str, PregelProtocol]]:
         for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse):
             yield name, node
 

libs/langgraph/langgraph/pregel/algo.py

@@ -630,6 +630,7 @@ def prepare_single_task(
                     task_id,
                     task_path[:3],
                     writers=proc.flat_writers,
+                    subgraphs=proc.subgraphs,
                 )
         else:
             return PregelTask(task_id, packet.node, task_path[:3])
@@ -754,6 +755,7 @@ def prepare_single_task(
                         task_id,
                         task_path[:3],
                         writers=proc.flat_writers,
+                        subgraphs=proc.subgraphs,
                     )
             else:
                 return PregelTask(task_id, name, task_path[:3])

libs/langgraph/langgraph/pregel/debug.py

@@ -30,7 +30,6 @@ from langgraph.constants import (
     TAG_HIDDEN,
 )
 from langgraph.pregel.io import read_channels
-from langgraph.pregel.utils import find_subgraph_pregel
 from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot
 from langgraph.utils.config import patch_checkpoint_map
 
@@ -157,7 +156,7 @@ def map_debug_checkpoint(
     task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {}
 
     for task in tasks:
-        if not find_subgraph_pregel(task.proc):
+        if not task.subgraphs:
             continue
 
         # assemble checkpoint_ns for this task

libs/langgraph/langgraph/pregel/read.py

@@ -22,7 +22,9 @@ from langchain_core.runnables.base import Input, Other, coerce_to_runnable
 from langchain_core.runnables.utils import ConfigurableFieldSpec
 
 from langgraph.constants import CONF, CONFIG_KEY_READ
+from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.retry import RetryPolicy
+from langgraph.pregel.utils import find_subgraph_pregel
 from langgraph.pregel.write import ChannelWrite
 from langgraph.utils.config import merge_configs
 from langgraph.utils.runnable import RunnableCallable, RunnableSeq
@@ -145,6 +147,9 @@ class PregelNode(Runnable):
     metadata: Optional[Mapping[str, Any]]
     """Metadata to attach to the node for tracing."""
 
+    subgraphs: Sequence[PregelProtocol]
+    """Subgraphs used by the node."""
+
     def __init__(
         self,
         *,
@@ -165,9 +170,21 @@ class PregelNode(Runnable):
         self.retry_policy = retry_policy
         self.tags = tags
         self.metadata = metadata
+        if self.bound is not DEFAULT_BOUND:
+            try:
+                subgraph = find_subgraph_pregel(self.bound)
+            except Exception:
+                subgraph = None
+            if subgraph:
+                self.subgraphs = [subgraph]
+            else:
+                self.subgraphs = []
+        else:
+            self.subgraphs = []
 
     def copy(self, update: dict[str, Any]) -> PregelNode:
         attrs = {**self.__dict__, **update}
+        attrs.pop("subgraphs")
         return PregelNode(**attrs)
 
     @cached_property
@@ -205,12 +222,12 @@ class PregelNode(Runnable):
             return self.bound
 
     def join(self, channels: Sequence[str]) -> PregelNode:
-        assert isinstance(channels, list) or isinstance(
-            channels, tuple
-        ), "channels must be a list or tuple"
-        assert isinstance(
-            self.channels, dict
-        ), "all channels must be named when using .join()"
+        assert isinstance(channels, list) or isinstance(channels, tuple), (
+            "channels must be a list or tuple"
+        )
+        assert isinstance(self.channels, dict), (
+            "all channels must be named when using .join()"
+        )
         return self.copy(
             update=dict(
                 channels={

libs/langgraph/langgraph/pregel/utils.py

@@ -26,7 +26,7 @@ def get_new_channel_versions(
     return new_versions
 
 
-def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]:
+def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]:
     from langgraph.pregel import Pregel
 
     candidates: list[Runnable] = [candidate]

libs/langgraph/langgraph/types.py

@@ -25,6 +25,7 @@ from typing_extensions import Self
 from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
 
 if TYPE_CHECKING:
+    from langgraph.pregel.protocol import PregelProtocol
     from langgraph.store.base import BaseStore
 
 
@@ -155,6 +156,7 @@ class PregelExecutableTask(NamedTuple):
     path: tuple[Union[str, int, tuple], ...]
     scheduled: bool = False
     writers: Sequence[Runnable] = ()
+    subgraphs: Sequence["PregelProtocol"] = ()
 
 
 class StateSnapshot(NamedTuple):

libs/langgraph/tests/test_pregel.py

@@ -6256,6 +6256,8 @@ def test_merging_updates_command_parent():
         ),
         ((), {"node_3": {"bar": ["node_3"]}}),
     ]
+
+
 def test_entrypoint_output_schema_with_return_and_save() -> None:
     """Test output schema inference with entrypoint.final."""