From 487fc8a8ca2983c134f26698d10933e5f035c96f Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 30 Aug 2024 15:13:56 -0700
Subject: [PATCH] Detect subgraphs called inside function nodes

---
 libs/langgraph/langgraph/graph/graph.py     |  8 ++++--
 libs/langgraph/langgraph/pregel/__init__.py | 30 ++++++++++++++++-----
 libs/langgraph/tests/test_pregel.py         |  5 +++-
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py
index 04c27ebaa..d466a3c96 100644
--- a/libs/langgraph/langgraph/graph/graph.py
+++ b/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()
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 8419eb913..01f0ad91d 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/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
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 31484cfe3..2f85540f1 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/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)