From 5444563fdcbf93e7fe34153dd0ba10dc77ab719e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 15 Jul 2024 11:26:05 -0700 Subject: [PATCH] Only render end node if there is an edge leading to it (#1022) --- libs/langgraph/langgraph/graph/graph.py | 23 ++- .../tests/__snapshots__/test_pregel.ambr | 168 +++++++++--------- 2 files changed, 99 insertions(+), 92 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 89f5c4f32..7dbbf184c 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -465,9 +465,16 @@ class CompiledGraph(Pregel): start_nodes: dict[str, DrawableNode] = { START: graph.add_node(self.get_input_schema(config), START) } - end_nodes: dict[str, DrawableNode] = { - END: graph.add_node(self.get_output_schema(config), END) - } + end_nodes: dict[str, DrawableNode] = {} + + def add_edge( + start: str, end: str, label: Optional[str] = None, conditional: bool = False + ) -> None: + if end == END and END not in end_nodes: + end_nodes[END] = graph.add_node(self.get_output_schema(config), END) + return graph.add_edge( + start_nodes[start], end_nodes[end], label, conditional + ) for key, (node, metadata) in self.builder.nodes.items(): if xray: @@ -494,7 +501,7 @@ class CompiledGraph(Pregel): start_nodes[key] = n end_nodes[key] = n for start, end in sorted(self.builder._all_edges): - graph.add_edge(start_nodes[start], end_nodes[end]) + add_edge(start, end) for start, branches in self.builder.branches.items(): default_ends = { **{k: k for k in self.builder.nodes if k != start}, @@ -508,13 +515,13 @@ class CompiledGraph(Pregel): else: ends = default_ends for label, end in ends.items(): - graph.add_edge( - start_nodes[start], - end_nodes[end], + add_edge( + start, + end, label if label != end else None, conditional=True, ) if branch.then is not None: - graph.add_edge(start_nodes[end], end_nodes[branch.then]) + add_edge(end, branch.then) return graph diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index ae3c2439f..bb751932f 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -16,11 +16,11 @@ %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; __start__([__start__]):::first - __end__([__end__]):::last prepare(prepare) tool_two_slow(tool_two_slow) tool_two_fast(tool_two_fast) finish(finish) + __end__([__end__]):::last __start__ --> prepare; finish --> __end__; prepare -.-> tool_two_slow; @@ -48,11 +48,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "left", "type": "runnable", @@ -76,6 +71,11 @@ ], "name": "right" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -129,11 +129,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "left", "type": "runnable", @@ -157,6 +152,11 @@ ], "name": "right" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -210,11 +210,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "get_weather", "type": "runnable", @@ -226,6 +221,11 @@ ], "name": "get_weather" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -265,11 +265,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "agent", "type": "runnable", @@ -298,6 +293,11 @@ "version": 2, "variant": "b" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -340,10 +340,10 @@ %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; __start__([__start__]):::first - __end__([__end__]):::last agent(agent) tools(tools
version = 2 variant = b) + __end__([__end__]):::last __start__ --> agent; tools --> agent; agent -.  continue  .-> tools; @@ -364,22 +364,17 @@ "data": "__start__" }, { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": 2, + "id": 1, "type": "schema", "data": "ParallelInput" }, { - "id": 3, + "id": 2, "type": "schema", "data": "ParallelOutput" }, { - "id": 4, + "id": 3, "type": "runnable", "data": { "id": [ @@ -392,7 +387,7 @@ } }, { - "id": 5, + "id": 4, "type": "runnable", "data": { "id": [ @@ -405,7 +400,7 @@ } }, { - "id": 6, + "id": 5, "type": "runnable", "data": { "id": [ @@ -418,7 +413,7 @@ } }, { - "id": 7, + "id": 6, "type": "runnable", "data": { "id": [ @@ -441,49 +436,54 @@ ], "name": "tools" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ + { + "source": 3, + "target": 4 + }, { "source": 4, "target": 5 }, + { + "source": 1, + "target": 3 + }, { "source": 5, + "target": 2 + }, + { + "source": 1, "target": 6 }, - { - "source": 2, - "target": 4 - }, { "source": 6, - "target": 3 - }, - { - "source": 2, - "target": 7 - }, - { - "source": 7, - "target": 3 + "target": 2 }, { "source": "__start__", - "target": 2 + "target": 1 }, { "source": "tools", - "target": 2 + "target": 1 }, { - "source": 3, + "source": 2, "target": "tools", "data": "continue", "conditional": true }, { - "source": 3, + "source": 2, "target": "__end__", "data": "exit", "conditional": true @@ -523,11 +523,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "agent", "type": "runnable", @@ -552,6 +547,11 @@ ], "name": "tools" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -643,11 +643,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "agent", "type": "runnable", @@ -672,6 +667,11 @@ ], "name": "tools" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -723,9 +723,9 @@ %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; __start__([__start__]):::first - __end__([__end__]):::last inner(inner) side(side) + __end__([__end__]):::last __start__ --> inner; inner --> side; side --> __end__; @@ -790,11 +790,6 @@ 'id': '__start__', 'type': 'schema', }), - dict({ - 'data': '__end__', - 'id': '__end__', - 'type': 'schema', - }), dict({ 'data': dict({ 'id': list([ @@ -812,11 +807,6 @@ 'id': 'tool_two:__start__', 'type': 'schema', }), - dict({ - 'data': 'tool_two:__end__', - 'id': 'tool_two:__end__', - 'type': 'schema', - }), dict({ 'data': dict({ 'id': list([ @@ -841,6 +831,11 @@ 'id': 'tool_two:tool_two_fast', 'type': 'runnable', }), + dict({ + 'data': 'tool_two:__end__', + 'id': 'tool_two:__end__', + 'type': 'schema', + }), dict({ 'data': dict({ 'id': list([ @@ -853,6 +848,11 @@ 'id': 'tool_three', 'type': 'runnable', }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), ]), }) # --- @@ -861,13 +861,13 @@ %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; __start__([__start__]):::first - __end__([__end__]):::last tool_one(tool_one) tool_two___start__(__start__) - tool_two___end__(__end__) tool_two_tool_two_slow(tool_two_slow) tool_two_tool_two_fast(tool_two_fast) + tool_two___end__(__end__) tool_three(tool_three) + __end__([__end__]):::last subgraph tool_two tool_two___start__ -.-> tool_two_tool_two_slow; tool_two_tool_two_slow --> tool_two___end__; @@ -901,11 +901,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "agent", "type": "runnable", @@ -931,6 +926,11 @@ ], "name": "tools" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -983,11 +983,6 @@ "type": "schema", "data": "__start__" }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, { "id": "agent", "type": "runnable", @@ -1013,6 +1008,11 @@ ], "name": "tools" } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" } ], "edges": [ @@ -1082,9 +1082,9 @@ %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; __start__([__start__]):::first - __end__([__end__]):::last tool_two_slow(tool_two_slow) tool_two_fast(tool_two_fast) + __end__([__end__]):::last __start__ -.-> tool_two_slow; tool_two_slow --> __end__; __start__ -.-> tool_two_fast;