Merge pull request #336 from langchain-ai/nc/22apr/fix-repeat-conditions

Fix issue when drawing graphs w repeat conditions

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)

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
   '''
       +-----------+  

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