diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 173798b3d..89004d23c 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -1,5 +1,5 @@ import logging -from collections import defaultdict +from collections import Counter, defaultdict from typing import ( Any, Awaitable, @@ -402,6 +402,9 @@ class CompiledGraph(Pregel): end_nodes[key] = n for start, end in sorted(self.graph._all_edges): graph.add_edge(start_nodes[start], end_nodes[end]) + branches_by_name = Counter( + name for _, branches in self.graph.branches.items() for name in branches + ) for start, branches in self.graph.branches.items(): for name, branch in branches.items(): ends = branch.ends or { @@ -410,7 +413,10 @@ class CompiledGraph(Pregel): } if add_condition_nodes is True: - cond = graph.add_node(branch.condition, name) + cond = graph.add_node( + branch.condition, + f"{start}_{name}" if branches_by_name[name] > 1 else name, + ) graph.add_edge(start_nodes[start], cond) for label, end in ends.items(): graph.add_edge(cond, end_nodes[end], label, conditional=True) diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index eca8025a3..7d9937b26 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1884,6 +1884,24 @@ +---------+ +--------+ ''' # --- +# name: test_repeat_condition + ''' + graph TD; + __start__ --> Researcher; + Researcher --> Researcher_router; + Researcher_router -. continue .-> Chart_Generator; + Researcher_router -. call_tool .-> Call_Tool; + Researcher_router -. end .-> __end__; + Chart_Generator --> Chart_Generator_router; + Chart_Generator_router -. continue .-> Researcher; + Chart_Generator_router -. call_tool .-> Call_Tool; + Chart_Generator_router -. end .-> __end__; + Call_Tool --> condition; + condition -. Researcher .-> Researcher; + condition -. Chart Generator .-> Chart_Generator; + + ''' +# --- # name: test_simple_multi_edge ''' +-----------+ diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 25accdf50..1ca671249 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -3534,3 +3534,42 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert app.invoke({"my_key": "my value"}) == { "my_key": "my value there and back again" } + + +def test_repeat_condition(snapshot: SnapshotAssertion) -> None: + class AgentState(TypedDict): + hello: str + + def router(state: AgentState) -> str: + return "hmm" + + workflow = StateGraph(AgentState) + workflow.add_node("Researcher", lambda x: x) + workflow.add_node("Chart Generator", lambda x: x) + workflow.add_node("Call Tool", lambda x: x) + workflow.add_conditional_edges( + "Researcher", + router, + {"continue": "Chart Generator", "call_tool": "Call Tool", "end": END}, + ) + workflow.add_conditional_edges( + "Chart Generator", + router, + {"continue": "Researcher", "call_tool": "Call Tool", "end": END}, + ) + workflow.add_conditional_edges( + "Call Tool", + # Each agent node updates the 'sender' field + # the tool calling node does not, meaning + # this edge will route back to the original agent + # who invoked the tool + lambda x: x["sender"], + { + "Researcher": "Researcher", + "Chart Generator": "Chart Generator", + }, + ) + workflow.set_entry_point("Researcher") + + app = workflow.compile() + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot