WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 01:37:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

Improve logic for drawing virtual end node

Browse Source 
- Don't draw unexpected edges to END node
- If a conditional edge has a custom label for END node then draw it
This commit is contained in:
Nuno Campos
2025-04-25 14:37:53 -07:00
parent 746b5f0730
commit 34d591dedf
6 changed files with 234 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/branch.py
+3 -2
View File
@@ -142,8 +142,8 @@ class Branch(NamedTuple):
),
list(
zip_longest(
writer([e for e in self.ends.values() if e != END]),
[str(la) for la, e in self.ends.items() if e != END],
writer([e for e in self.ends.values()]),
[str(la) for la, e in self.ends.items()],
)
)
if self.ends
@@ -211,6 +211,7 @@ class Branch(NamedTuple):
]
else:
destinations = cast(Sequence[Union[Send, str]], result)
destinations = [d for d in destinations if d != END]
if any(dest is None or dest == START for dest in destinations):
raise ValueError("Branch did not return a valid destination")
if any(p.node == END for p in destinations if isinstance(p, Send)):
libs/langgraph/langgraph/graph/state.py
+21 -28
View File
@@ -845,29 +845,26 @@ class CompiledStateGraph(CompiledGraph):
def get_writes(
packets: Sequence[Union[str, Send]],
) -> Sequence[Union[ChannelWriteEntry, Send]]:
if filtered := [p for p in packets if p != END]:
writes = [
(
ChannelWriteEntry(CHANNEL_BRANCH_TO.format(p), None)
if not isinstance(p, Send)
else p
writes = [
(
ChannelWriteEntry(CHANNEL_BRANCH_TO.format(p), None)
if not isinstance(p, Send)
else p
)
for p in packets
]
if branch.then and branch.then != END:
writes.append(
ChannelWriteEntry(
f"branch:{start}:{name}::then",
WaitForNames(
frozenset(
p.node if isinstance(p, Send) else p for p in packets
)
),
)
for p in filtered
]
if branch.then and branch.then != END:
writes.append(
ChannelWriteEntry(
f"branch:{start}:{name}::then",
WaitForNames(
frozenset(
p.node if isinstance(p, Send) else p
for p in filtered
)
),
)
)
return writes
return []
)
return writes
if with_reader:
# get schema
@@ -1067,13 +1064,9 @@ def _control_static(
ends: Union[tuple[str, ...], dict[str, str]],
) -> Sequence[tuple[str, Any, Optional[str]]]:
if isinstance(ends, dict):
return [
(CHANNEL_BRANCH_TO.format(k), None, label)
for k, label in ends.items()
if k != END
]
return [(CHANNEL_BRANCH_TO.format(k), None, label) for k, label in ends.items()]
else:
return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends if e != END]
return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends]
def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
libs/langgraph/langgraph/pregel/draw.py
+38 -11
View File
@@ -110,6 +110,11 @@ def draw_graph(
static_seen.add(w)
# apply static writes
if writes := ChannelWrite.get_static_writes(w):
# END writes are not written, but become edges directly
for w in writes:
if w[0] == END:
edges.add((task.name, w[0], True, w[2]))
writes = [t for t in writes if t[0] != END]
conditionals.update(
{(task.name, *t[:2]): t[2] for t in writes}
)
@@ -175,31 +180,36 @@ def draw_graph(
if START not in nodes:
graph.add_node(None, START)
for task in start_tasks.values():
graph.add_edge(graph.nodes[START], graph.nodes[task.name])
add_edge(graph, START, task.name)
# add discovered edges
for src, dest, is_conditional, label in sorted(edges):
graph.add_edge(
graph.nodes[src],
graph.nodes[dest],
add_edge(
graph,
src,
dest,
data=label if label != dest else None,
conditional=is_conditional,
)
# add end edges
if step_sources:
end = graph.add_node(None, END)
termini = {d for _, d, _, _ in edges}.difference(s for s, _, _, _ in edges)
for src in sorted(termini.union(step_sources)):
graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
termini = {d for _, d, _, _ in edges if d != END}.difference(
s for s, _, _, _ in edges
)
if termini:
for src in sorted(termini):
add_edge(graph, src, END)
elif len(step_sources) == 1:
for src in sorted(step_sources):
add_edge(graph, src, END, conditional=True)
# replace subgraphs
for name, subgraph in subgraphs.items():
subgraph.trim_first_node()
subgraph.trim_last_node()
if (
len(subgraph.nodes) > 1
and name in graph.nodes
and subgraph.first_node()
and subgraph.last_node()
):
subgraph.trim_first_node()
subgraph.trim_last_node()
# replace the node with the subgraph
graph.nodes.pop(name)
first, last = graph.extend(subgraph, prefix=name)
@@ -210,3 +220,20 @@ def draw_graph(
graph.edges[idx] = edge.copy(target=cast(Node, first).id)
return graph
def add_edge(
graph: Graph,
source: str,
target: str,
*,
data: Optional[Any] = None,
conditional: bool = False,
) -> None:
"""Add an edge to the graph."""
for edge in graph.edges:
if edge.source == source and edge.target == target:
return
if target not in graph.nodes and target == END:
graph.add_node(None, END)
graph.add_edge(graph.nodes[source], graph.nodes[target], data, conditional)
libs/langgraph/tests/__snapshots__/test_large_cases.ambr
+8 -7
View File
@@ -86,6 +86,12 @@
"source": "__start__",
"target": "agent"
},
{
"source": "agent",
"target": "__end__",
"data": "exit",
"conditional": true
},
{
"source": "agent",
"target": "tools",
@@ -95,11 +101,6 @@
{
"source": "tools",
"target": "agent"
},
{
"source": "agent",
"target": "__end__",
"conditional": true
}
]
}
@@ -109,9 +110,9 @@
'''
graph TD;
__start__ --> agent;
agent -.  exit  .-> __end__;
agent -.  continue  .-> tools;
tools --> agent;
agent -.-> __end__;
'''
# ---
@@ -130,9 +131,9 @@
__start__([<p>__start__</p>]):::first
__end__([<p>__end__</p>]):::last
__start__ --> agent;
agent -. &nbsp;exit&nbsp; .-> __end__;
agent -. &nbsp;continue&nbsp; .-> tools;
tools --> agent;
agent -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+119 -26
View File
@@ -67,7 +67,8 @@
},
{
"source": "left",
"target": "__end__"
"target": "__end__",
"conditional": true
},
{
"source": "right",
@@ -82,7 +83,7 @@
graph TD;
__start__ -. &nbsp;go-left&nbsp; .-> left;
__start__ -. &nbsp;go-right&nbsp; .-> right;
left --> __end__;
left -.-> __end__;
right --> __end__;
'''
@@ -481,6 +482,17 @@
'''
# ---
# name: test_migration_graph
'''
graph TD;
B -. &nbsp;X&nbsp; .-> C;
B -. &nbsp;Y&nbsp; .-> D;
D --> B;
__start__ --> B;
C --> __end__;
'''
# ---
# name: test_multiple_sinks_subgraphs
'''
---
@@ -532,12 +544,14 @@
---
graph TD;
__start__([<p>__start__</p>]):::first
inner(inner)
side(side)
__end__([<p>__end__</p>]):::last
__start__ --> inner;
inner --> side;
__start__ --> inner_up;
inner_up --> side;
side --> __end__;
subgraph inner
inner_up(up)
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
@@ -550,8 +564,48 @@
dict({
'conditional': True,
'source': '__start__',
'target': 'tool_one',
}),
dict({
'conditional': True,
'source': '__start__',
'target': 'tool_three',
}),
dict({
'conditional': True,
'source': '__start__',
'target': 'tool_two:__start__',
}),
dict({
'source': 'tool_one',
'target': '__end__',
}),
dict({
'source': 'tool_three',
'target': '__end__',
}),
dict({
'source': 'tool_two:__end__',
'target': '__end__',
}),
dict({
'conditional': True,
'source': 'tool_two:__start__',
'target': 'tool_two:tool_two_fast',
}),
dict({
'conditional': True,
'source': 'tool_two:__start__',
'target': 'tool_two:tool_two_slow',
}),
dict({
'source': 'tool_two:tool_two_fast',
'target': 'tool_two:__end__',
}),
dict({
'source': 'tool_two:tool_two_slow',
'target': 'tool_two:__end__',
}),
]),
'nodes': list([
dict({
@@ -580,19 +634,6 @@
'id': 'tool_one',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'graph',
'state',
'CompiledStateGraph',
]),
'name': 'tool_two',
}),
'id': 'tool_two',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
@@ -609,6 +650,48 @@
dict({
'id': '__end__',
}),
dict({
'data': dict({
'id': list([
'langchain',
'schema',
'runnable',
'RunnablePassthrough',
]),
'name': 'tool_two:__start__',
}),
'id': 'tool_two:__start__',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_two:tool_two_slow',
}),
'id': 'tool_two:tool_two_slow',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_two:tool_two_fast',
}),
'id': 'tool_two:tool_two_fast',
'type': 'runnable',
}),
dict({
'id': 'tool_two:__end__',
}),
]),
})
# ---
@@ -620,12 +703,26 @@
curve: linear
---
graph TD;
__start__(<p>__start__</p>)
__start__([<p>__start__</p>]):::first
tool_one(tool_one)
tool_two(tool_two)
tool_three(tool_three)
__end__(<p>__end__</p>)
__start__ -.-> __end__;
__end__([<p>__end__</p>]):::last
__start__ -.-> tool_one;
__start__ -.-> tool_three;
__start__ -.-> tool_two___start__;
tool_one --> __end__;
tool_three --> __end__;
tool_two___end__ --> __end__;
subgraph tool_two
tool_two___start__(<p>__start__</p>)
tool_two_tool_two_slow(tool_two_slow)
tool_two_tool_two_fast(tool_two_fast)
tool_two___end__(<p>__end__</p>)
tool_two___start__ -.-> tool_two_tool_two_fast;
tool_two___start__ -.-> tool_two_tool_two_slow;
tool_two_tool_two_fast --> tool_two___end__;
tool_two_tool_two_slow --> tool_two___end__;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
@@ -642,9 +739,6 @@
Researcher -. &nbsp;call_tool&nbsp; .-> Call_Tool;
Researcher -. &nbsp;continue&nbsp; .-> Chart_Generator;
__start__ --> Researcher;
Call_Tool -.-> __end__;
Chart_Generator -.-> __end__;
Researcher -.-> __end__;
Researcher -. &nbsp;redo&nbsp; .-> Researcher;
'''
@@ -1014,7 +1108,6 @@
dict({
'id': 'conduct_interview:__end__',
}),
]),
})
# ---
libs/langgraph/tests/test_pregel.py
+45 -1
View File
@@ -3491,6 +3491,7 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
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",
["tool_two_slow", "tool_two_fast"],
then=END,
)
tool_two = tool_two_graph.compile()
@@ -3499,7 +3500,9 @@ def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
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)
graph.set_conditional_entry_point(
lambda s: "tool_one", ["tool_one", "tool_two", "tool_three"], then=END
)
app = graph.compile()
assert app.get_graph(xray=True).to_json() == snapshot
@@ -8120,3 +8123,44 @@ def test_batch_update_as_input(
]
assert new_history == history
def test_migration_graph(snapshot: SnapshotAssertion) -> None:
from pydantic import BaseModel
class DummyState(BaseModel):
pass_count: int = 0
def increment_pass_count(state: DummyState):
state.pass_count += 1
return state
def route_b(state: DummyState):
if state.pass_count == 0:
return "X"
else:
return "Y"
migration_graph = StateGraph(DummyState)
migration_graph.add_node("B", increment_pass_count)
migration_graph.add_node("C", increment_pass_count)
migration_graph.add_node("D", increment_pass_count)
migration_graph.add_edge(START, "B")
migration_graph.add_conditional_edges(
"B",
route_b,
{
"X": "C",
"Y": "D",
},
)
migration_graph.add_edge("D", "B")
migration_graph.add_edge("C", END)
app = migration_graph.compile()
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot