From 42b13ef5769c24ac6a53ad26d62de331cebb15e8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 24 Apr 2024 10:08:23 -0700 Subject: [PATCH] Improvements to get_graph(xray=) - Support xray: int, to control depth, eg xray=1 only exposes one level deep - In mermaid draw a containing box around subgraphs created from xray --- langgraph/graph/graph.py | 22 +-- langgraph/utils.py | 32 ++++ poetry.lock | 8 +- pyproject.toml | 2 +- tests/__snapshots__/test_pregel.ambr | 249 +++++++++++++++++++++++++++ tests/test_pregel.py | 30 ++++ 6 files changed, 327 insertions(+), 16 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 8150fa3f8..1c5f9c957 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -15,9 +15,6 @@ from typing import ( from langchain_core.runnables import Runnable from langchain_core.runnables.base import RunnableLike, coerce_to_runnable from langchain_core.runnables.config import RunnableConfig -from langchain_core.runnables.graph import ( - Graph as RunnableGraph, -) from langchain_core.runnables.graph import ( Node as RunnableGraphNode, ) @@ -28,7 +25,7 @@ from langgraph.constants import TAG_HIDDEN from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.utils import RunnableCallable +from langgraph.utils import DrawableGraph, RunnableCallable logger = logging.getLogger(__name__) @@ -380,11 +377,11 @@ class CompiledGraph(Pregel): self, config: Optional[RunnableConfig] = None, *, - xray: bool = False, + xray: Union[int, bool] = False, add_condition_nodes: bool = True, - ) -> RunnableGraph: + ) -> DrawableGraph: """Returns a drawable representation of the computation graph.""" - graph = RunnableGraph() + graph = DrawableGraph() start_nodes: dict[str, RunnableGraphNode] = { START: graph.add_node(self.get_input_schema(config), START) } @@ -395,16 +392,19 @@ class CompiledGraph(Pregel): for key, node in self.graph.nodes.items(): if xray: subgraph = ( - node.get_graph(config=config, xray=xray) + node.get_graph( + config=config, + xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray, + ) if isinstance(node, CompiledGraph) else node.get_graph(config=config) ) subgraph.trim_first_node() subgraph.trim_last_node() if len(subgraph.nodes) > 1: - graph.extend(subgraph) - start_nodes[key] = subgraph.last_node() - end_nodes[key] = subgraph.first_node() + end_nodes[key], start_nodes[key] = graph.extend( + subgraph, prefix=key + ) else: n = graph.add_node(node, key) start_nodes[key] = n diff --git a/langgraph/utils.py b/langgraph/utils.py index 0604fd81e..fbbd61f17 100644 --- a/langgraph/utils.py +++ b/langgraph/utils.py @@ -3,6 +3,7 @@ from typing import Any, Awaitable, Callable, Optional from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.config import merge_configs +from langchain_core.runnables.graph import Edge, Graph, Node, is_uuid # Before Python 3.11 native StrEnum is not available @@ -67,3 +68,34 @@ class RunnableCallable(Runnable): if isinstance(ret, Runnable) and self.recurse: return await ret.ainvoke(input, config) return ret + + +class DrawableGraph(Graph): + def extend( + self, graph: Graph, prefix: str = "" + ) -> tuple[Optional[Node], Optional[Node]]: + if all(is_uuid(node.id) for node in graph.nodes.values()): + super().extend(graph) + return graph.first_node(), graph.last_node() + + new_nodes = { + f"{prefix}:{k}": Node(f"{prefix}:{k}", v.data) + for k, v in graph.nodes.items() + } + new_edges = [ + Edge( + f"{prefix}:{edge.source}", + f"{prefix}:{edge.target}", + edge.data, + edge.conditional, + ) + for edge in graph.edges + ] + self.nodes.update(new_nodes) + self.edges.extend(new_edges) + first = graph.first_node() + last = graph.last_node() + return ( + Node(f"{prefix}:{first.id}", first.data) if first else None, + Node(f"{prefix}:{last.id}", last.data) if last else None, + ) diff --git a/poetry.lock b/poetry.lock index cb4b6cb5d..98d279991 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1710,13 +1710,13 @@ extended-testing = ["aiosqlite (>=0.19.0,<0.20.0)", "aleph-alpha-client (>=2.15. [[package]] name = "langchain-core" -version = "0.1.45" +version = "0.1.46" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langchain_core-0.1.45-py3-none-any.whl", hash = "sha256:91eff20de0bcf5f025e1d8c4582cb597a9c17527965eb03b314486e7c834e7df"}, - {file = "langchain_core-0.1.45.tar.gz", hash = "sha256:526532c1af279a9e2debe7a4e143ba6e980cf90b5ab2e0991c2230ee04c693e2"}, + {file = "langchain_core-0.1.46-py3-none-any.whl", hash = "sha256:1c0befcd2665dd4aa153318aa9bf729071644b4c179e491769b8e583b4bf7441"}, + {file = "langchain_core-0.1.46.tar.gz", hash = "sha256:17c416349f5c7a9808e70e3725749a3a2df5088f1ecca045c883871aa95f9c9e"}, ] [package.dependencies] @@ -4094,4 +4094,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "5d0e5b014e355f65a1731548531a53313cff61d729312d64f1f3b339673e2f5c" +content-hash = "4c5fea39af2c255ddbc99a3b0ccbdb94e53c3526e210f023f71c95b1d34c0a98" diff --git a/pyproject.toml b/pyproject.toml index 031d5baa1..ffba91ca5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph" [tool.poetry.dependencies] python = ">=3.9.0,<4.0" -langchain-core = "^0.1.42" +langchain-core = "^0.1.46" [tool.poetry.group.test.dependencies] diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 4d4a0df6e..85d47dd11 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1901,6 +1901,255 @@ ''' # --- +# name: test_nested_graph_xray + dict({ + 'edges': list([ + dict({ + 'conditional': True, + 'data': 'tool_two_slow', + 'source': 'tool_two:condition', + 'target': 'tool_two:tool_two_slow', + }), + dict({ + 'source': 'tool_two:tool_two_slow', + 'target': 'tool_two:__end__', + }), + dict({ + 'conditional': True, + 'data': 'tool_two_fast', + 'source': 'tool_two:condition', + 'target': 'tool_two:tool_two_fast', + }), + dict({ + 'source': 'tool_two:tool_two_fast', + 'target': 'tool_two:__end__', + }), + dict({ + 'source': '__start__', + 'target': 'condition', + }), + dict({ + 'conditional': True, + 'data': 'tool_one', + 'source': 'condition', + 'target': 'tool_one', + }), + dict({ + 'source': 'tool_one', + 'target': '__end__', + }), + dict({ + 'conditional': True, + 'data': 'tool_two', + 'source': 'condition', + 'target': 'tool_two:condition', + }), + dict({ + 'source': 'tool_two:__end__', + 'target': '__end__', + }), + dict({ + 'conditional': True, + 'data': 'tool_three', + 'source': 'condition', + 'target': 'tool_three', + }), + dict({ + 'source': 'tool_three', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain_core', + 'runnables', + 'base', + 'RunnableLambda', + ]), + 'name': 'logic', + }), + 'id': 'tool_one', + 'type': 'runnable', + }), + dict({ + 'data': 'tool_two:__end__', + 'id': 'tool_two:__end__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain_core', + 'runnables', + 'base', + 'RunnableLambda', + ]), + 'name': 'logic', + }), + 'id': 'tool_two:tool_two_slow', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain_core', + 'runnables', + 'base', + 'RunnableLambda', + ]), + 'name': 'logic', + }), + 'id': 'tool_two:tool_two_fast', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain_core', + 'runnables', + 'base', + 'RunnableLambda', + ]), + 'name': 'RunnableLambda', + }), + 'id': 'tool_two:condition', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain_core', + 'runnables', + 'base', + 'RunnableLambda', + ]), + 'name': 'logic', + }), + 'id': 'tool_three', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain_core', + 'runnables', + 'base', + 'RunnableLambda', + ]), + 'name': 'RunnableLambda', + }), + 'id': 'condition', + 'type': 'runnable', + }), + ]), + }) +# --- +# name: test_nested_graph_xray.1 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-----------+ + | condition | + ..+-----------+... + ..... . ..... + ... . ... + ... . ... + +----------+ +----------+ +------------+ + | tool_one |* | tool_two | | tool_three | + +----------+ *** +----------+ **+------------+ + ***** * ***** + *** * *** + *** * *** + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_nested_graph_xray.2 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-----------+ + | condition |..... + .....+-----------+... ....... + .... .... ......... + ..... ... ....... + ... .... ......... + +--------------------+ .. .... + | tool_two:condition | . . + +--------------------+ . . + ... ... . . + ... ... . . + .. .. . . + +------------------------+ +------------------------+ . . + | tool_two:tool_two_slow | | tool_two:tool_two_fast | . . + +------------------------+ +------------------------+ . . + *** *** . . + *** *** . . + ** ** . . + +------------------+ +----------+ +------------+ + | tool_two:__end__ | | tool_one | ****| tool_three | + +------------------+** ***+----------+******** +------------+ + **** *** ******** + ***** ************* + *** ****** + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_nested_graph_xray.3 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__[__start__]:::startclass; + __end__[__end__]:::endclass; + tool_one([tool_one]):::otherclass; + tool_two___end__([__end__]):::otherclass; + tool_two_tool_two_slow([tool_two_slow]):::otherclass; + tool_two_tool_two_fast([tool_two_fast]):::otherclass; + tool_two_condition([condition]):::otherclass; + tool_three([tool_three]):::otherclass; + condition([condition]):::otherclass; + subgraph tool_two + tool_two_condition -. tool_two_slow .-> tool_two_tool_two_slow; + tool_two_tool_two_slow --> tool_two___end__; + tool_two_condition -. tool_two_fast .-> tool_two_tool_two_fast; + tool_two_tool_two_fast --> tool_two___end__; + end + __start__ --> condition; + condition -. tool_one .-> tool_one; + tool_one --> __end__; + condition -. tool_two .-> tool_two_condition; + tool_two___end__ --> __end__; + condition -. tool_three .-> tool_three; + tool_three --> __end__; + classDef startclass fill:#ffdfba; + classDef endclass fill:#baffc9; + classDef otherclass fill:#fad7de; + + ''' +# --- # name: test_prebuilt_chat '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' # --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index d184b3754..316d98817 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -3848,6 +3848,36 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None: assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"} +def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def logic(state: State): + pass + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two_slow", logic) + tool_two_graph.add_node("tool_two_fast", logic) + tool_two_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + then=END, + ) + tool_two = tool_two_graph.compile() + + graph = StateGraph(State) + graph.add_node("tool_one", logic) + graph.add_node("tool_two", tool_two) + graph.add_node("tool_three", logic) + graph.set_conditional_entry_point(lambda s: "tool_one", then=END) + app = graph.compile() + + assert app.get_graph(xray=True).to_json() == snapshot + assert app.get_graph().draw_ascii() == snapshot + assert app.get_graph(xray=True).draw_ascii() == snapshot + assert app.get_graph(xray=True).draw_mermaid() == snapshot + + def test_nested_graph(snapshot: SnapshotAssertion) -> None: def never_called_fn(state: Any): assert 0, "This function should never be called"