From fa467573d71b86361b8ced9f4075f586b1b831e7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 09:36:33 -0700 Subject: [PATCH 1/3] Enable xray for remote graphs --- libs/langgraph/langgraph/graph/graph.py | 72 ++++++++++++++++++++----- 1 file changed, 58 insertions(+), 14 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 137eb0862..9592dcfce 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,7 +32,9 @@ 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.remote import RemoteGraph from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.types import All, Checkpointer from langgraph.utils.runnable import RunnableLike, coerce_to_runnable @@ -418,7 +421,36 @@ class CompiledGraph(Pregel): *, xray: Union[int, bool] = False, ) -> DrawableGraph: - return self.get_graph(config, xray=xray) + """Returns a drawable representation of the computation graph.""" + # 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 - 1 + if isinstance(xray, int) and xray > 0 + else xray, + ) + for p in subpregels.values() + ) + ), + ) + } + else: + subgraphs = {} + + # draw the graph + return self._draw_graph(config, subgraphs=subgraphs) def get_graph( self, @@ -427,17 +459,34 @@ class CompiledGraph(Pregel): xray: Union[int, bool] = False, ) -> DrawableGraph: """Returns a drawable representation of the computation graph.""" + # gather subgraphs + if xray: + subgraphs = { + k: v.get_graph( + config, + xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray, + ) + 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 +512,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: From ddb29df6671b12eccf36d418bc4b7a3c7e1de18c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 09:39:23 -0700 Subject: [PATCH 2/3] Fix --- libs/langgraph/langgraph/graph/graph.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 9592dcfce..df81de4da 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -34,7 +34,6 @@ 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.remote import RemoteGraph from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.types import All, Checkpointer from langgraph.utils.runnable import RunnableLike, coerce_to_runnable @@ -422,6 +421,8 @@ 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: subpregels: dict[str, PregelProtocol] = { @@ -459,6 +460,8 @@ 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 = { From 5db1949ae378118fe0b8706ef74d725391731892 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 09:55:21 -0700 Subject: [PATCH 3/3] Fix --- libs/langgraph/langgraph/graph/graph.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index df81de4da..074364867 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -438,9 +438,9 @@ class CompiledGraph(Pregel): *( p.aget_graph( config, - xray=xray - 1 - if isinstance(xray, int) and xray > 0 - else xray, + xray=xray + if isinstance(xray, bool) or xray <= 0 + else xray - 1, ) for p in subpregels.values() ) @@ -467,7 +467,7 @@ class CompiledGraph(Pregel): subgraphs = { k: v.get_graph( config, - xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray, + xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1, ) for k, v in self.get_subgraphs() if isinstance(v, (CompiledGraph, RemoteGraph))