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;