diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 137eb0862..074364867 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -1,3 +1,4 @@ +import asyncio import logging from collections import defaultdict from typing import ( @@ -31,6 +32,7 @@ from langgraph.constants import ( ) from langgraph.graph.branch import Branch from langgraph.pregel import Channel, Pregel +from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.types import All, Checkpointer @@ -418,7 +420,38 @@ class CompiledGraph(Pregel): *, xray: Union[int, bool] = False, ) -> DrawableGraph: - return self.get_graph(config, xray=xray) + """Returns a drawable representation of the computation graph.""" + from langgraph.pregel.remote import RemoteGraph + + # gather subgraphs + if xray: + subpregels: dict[str, PregelProtocol] = { + k: v + async for k, v in self.aget_subgraphs() + if isinstance(v, (CompiledGraph, RemoteGraph)) + } + subgraphs = { + k: v + for k, v in zip( + subpregels, + await asyncio.gather( + *( + p.aget_graph( + config, + xray=xray + if isinstance(xray, bool) or xray <= 0 + else xray - 1, + ) + for p in subpregels.values() + ) + ), + ) + } + else: + subgraphs = {} + + # draw the graph + return self._draw_graph(config, subgraphs=subgraphs) def get_graph( self, @@ -427,17 +460,36 @@ class CompiledGraph(Pregel): xray: Union[int, bool] = False, ) -> DrawableGraph: """Returns a drawable representation of the computation graph.""" + from langgraph.pregel.remote import RemoteGraph + + # gather subgraphs + if xray: + subgraphs = { + k: v.get_graph( + config, + xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1, + ) + for k, v in self.get_subgraphs() + if isinstance(v, (CompiledGraph, RemoteGraph)) + } + else: + subgraphs = {} + + # draw the graph + return self._draw_graph(config, subgraphs=subgraphs) + + def _draw_graph( + self, + config: Optional[RunnableConfig] = None, + *, + subgraphs: dict[str, DrawableGraph] = {}, + ) -> DrawableGraph: + # create the graph graph = DrawableGraph() start_nodes: dict[str, DrawableNode] = { START: graph.add_node(self.get_input_schema(config), START) } end_nodes: dict[str, DrawableNode] = {} - if xray: - subgraphs = { - k: v for k, v in self.get_subgraphs() if isinstance(v, CompiledGraph) - } - else: - subgraphs = {} def add_edge( start: str, @@ -463,13 +515,8 @@ class CompiledGraph(Pregel): metadata["__interrupt"] = "before" elif key in self.interrupt_after_nodes: metadata["__interrupt"] = "after" - if xray and key in subgraphs: - subgraph = subgraphs[key].get_graph( - config=config, - xray=xray - 1 - if isinstance(xray, int) and not isinstance(xray, bool) and xray > 0 - else xray, - ) + if key in subgraphs: + subgraph = subgraphs[key] subgraph.trim_first_node() subgraph.trim_last_node() if len(subgraph.nodes) >= 1: