Detect subgraphs called inside function nodes

libs/langgraph/langgraph/graph/graph.py

@@ -484,6 +484,10 @@ class CompiledGraph(Pregel):
             START: graph.add_node(self.get_input_schema(config), START)
         }
         end_nodes: dict[str, DrawableNode] = {}
+        if xray:
+            subgraphs = dict(self.get_subgraphs())
+        else:
+            subgraphs = {}
 
         def add_edge(
             start: str, end: str, label: Optional[str] = None, conditional: bool = False
@@ -503,11 +507,11 @@ class CompiledGraph(Pregel):
                 metadata["__interrupt"] = "after"
             if xray:
                 subgraph = (
-                    node.get_graph(
+                    subgraphs[key].get_graph(
                         config=config,
                         xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray,
                     )
-                    if isinstance(node, CompiledGraph)
+                    if key in subgraphs
                     else node.get_graph(config=config)
                 )
                 subgraph.trim_first_node()

libs/langgraph/langgraph/pregel/__init__.py

@@ -28,6 +28,7 @@ from langchain_core.load.dump import dumpd
 from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
 from langchain_core.runnables import (
     Runnable,
+    RunnableLambda,
     RunnableSequence,
     RunnableSerializable,
 )
@@ -43,6 +44,7 @@ from langchain_core.runnables.config import (
 from langchain_core.runnables.utils import (
     ConfigurableFieldSpec,
     create_model,
+    get_function_nonlocals,
     get_unique_config_specs,
 )
 from langchain_core.tracers._streaming import _StreamingCallbackHandler
@@ -102,6 +104,7 @@ from langgraph.pregel.utils import (
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
+from langgraph.utils import RunnableCallable
 
 WriteValue = Union[
     Runnable[Input, Output],
@@ -356,13 +359,26 @@ class Pregel(
         for name, node in self.nodes.items():
             # find the subgraph, if any
             graph: Optional[Pregel] = None
-            if isinstance(node.bound, Pregel):
-                graph = node.bound
-            elif isinstance(node.bound, RunnableSequence):
-                for runnable in node.bound.steps:
-                    if isinstance(runnable, Pregel):
-                        graph = runnable
-                        break
+            candidates = [node.bound]
+            for candidate in candidates:
+                if isinstance(candidate, Pregel):
+                    graph = candidate
+                    break
+                elif isinstance(candidate, RunnableSequence):
+                    candidates.extend(candidate.steps)
+                elif isinstance(candidate, RunnableLambda):
+                    candidates.extend(candidate.deps)
+                elif isinstance(candidate, RunnableCallable):
+                    if candidate.func is not None:
+                        candidates.extend(
+                            nl.__self__ if hasattr(nl, "__self__") else nl
+                            for nl in get_function_nonlocals(candidate.func)
+                        )
+                    if candidate.afunc is not None:
+                        candidates.extend(
+                            nl.__self__ if hasattr(nl, "__self__") else nl
+                            for nl in get_function_nonlocals(candidate.afunc)
+                        )
             # if found, yield recursively
             if graph:
                 yield name, graph

libs/langgraph/tests/test_pregel.py

@@ -10209,10 +10209,13 @@ def test_weather_subgraph(
         else:
             return "normal_llm_node"
 
+    def weather_graph(state: RouterState):
+        return subgraph.invoke(state)
+
     graph = StateGraph(RouterState)
     graph.add_node(router_node)
     graph.add_node(normal_llm_node)
-    graph.add_node("weather_graph", subgraph)
+    graph.add_node("weather_graph", weather_graph)
     graph.add_edge(START, "router_node")
     graph.add_conditional_edges("router_node", route_after_prediction)
     graph.add_edge("normal_llm_node", END)