Merge pull request #53 from aurelio-labs/feat/mulitple-activated-branches-in-router

feat: allow duplicate node names in parallel

Author: jamescalam

graphai/graph.py

@@ -362,9 +362,19 @@ async def _execute_branch(
         callback: Callback,
         steps: int,
         stop_at_join: bool = False,
+        branch_id: str | None = None,
     ):
         """Recursively execute a branch starting from `current_node`.
-        When a node has multiple successors, run them concurrently and merge their outputs."""
+        When a node has multiple successors, run them concurrently and merge their outputs.
+
+        Args:
+            current_node: The node to start execution from
+            state: The current state dict to pass to the node
+            callback: The callback instance for events
+            steps: Current step count for max_steps limit
+            stop_at_join: If True, stop execution when reaching a join node
+            branch_id: Unique identifier for this branch (used for parallel duplicate nodes)
+        """
         while True:
             output = await self._invoke_node(current_node, state, callback)
             state = {**state, **output}  # merge node output into local state
@@ -396,19 +406,32 @@ async def _execute_branch(
             if len(next_nodes) == 1:
                 current_node = next_nodes[0]
             else:
-                # Run each branch concurrently
-                results = await asyncio.gather(
-                    *[
+                # Run each branch concurrently, assigning unique branch IDs
+                # This ensures that even duplicate nodes (same node appearing multiple times)
+                # are tracked as separate invocations
+                branch_tasks = []
+                for idx, n in enumerate(next_nodes):
+                    # Create unique branch ID combining parent branch, node name, and index
+                    new_branch_id = f"{branch_id or 'root'}_{n.name}_{idx}"
+                    branch_tasks.append(
                         self._execute_branch(
                             current_node=n,
                             state=state.copy(),
                             callback=callback,
                             steps=steps + 1,
                             stop_at_join=True,  # force parallel branches to wait at JoinEdge
+                            branch_id=new_branch_id,
                         )
-                        for n in next_nodes
-                    ]
+                    )
+
+                # Wait for ALL branches to complete (including duplicate node invocations)
+                logger.debug(
+                    f"Starting {len(branch_tasks)} parallel branches: "
+                    f"{[n.name for n in next_nodes]}"
                 )
+                results = await asyncio.gather(*branch_tasks)
+                logger.debug(f"All {len(results)} parallel branches completed")
+
                 # merge states returned by each branch
                 merged = state.copy()
                 for res in results:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "graphai-lib"
-version = "0.0.10"
+version = "0.0.11rc1"
 description = "Not an AI framework"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"

tests/unit/test_graph_parallel.py

@@ -370,3 +370,107 @@ async def end(input: dict):
     assert result["a_result"] == 1
     assert "b_result" not in result
     assert "parallel_results" not in result
+
+
+@pytest.mark.asyncio
+async def test_router_parallel_duplicate_node_names():
+    """Router returning duplicate node names should execute the same node multiple times.
+
+    Uses different sleep durations (1s vs 3s) to verify all branches are waited for.
+    """
+    import asyncio
+
+    call_count = {"tool_a": 0, "router": 0}
+    sleep_times = [3, 3, 1]  # Last invocation is fast, others are slow
+
+    @node(start=True)
+    async def start(input: dict):
+        return {}
+
+    @router(name="parallel_router")
+    async def parallel_router(input: dict):
+        call_count["router"] += 1
+        if call_count["router"] > 1:
+            return {"choice": "end"}
+        # Return the same node name multiple times
+        return {"choices": ["tool_a", "tool_a", "tool_a"]}
+
+    @node(name="tool_a")
+    async def tool_a(input: dict):
+        call_count["tool_a"] += 1
+        sleep_time = sleep_times.pop()
+        await asyncio.sleep(sleep_time)
+        return {"a_result": call_count["tool_a"]}
+
+    @node(end=True)
+    async def end(input: dict):
+        return {"final": "done"}
+
+    g = Graph()
+    g.add_node(start).add_node(parallel_router).add_node(tool_a).add_node(end)
+    g.add_edge(start, parallel_router)
+    g.add_edge(parallel_router, tool_a)
+    g.add_join([tool_a], parallel_router)
+    g.add_edge(parallel_router, end)
+
+    result = await g.execute({"input": {}})
+
+    # tool_a should have been called 3 times (once for each duplicate in choices)
+    assert call_count["tool_a"] == 3
+    # Router should have been called twice (first for parallel, second to continue to end)
+    assert call_count["router"] == 2
+    assert result["final"] == "done"
+
+
+@pytest.mark.asyncio
+async def test_router_parallel_mixed_duplicate_nodes():
+    """Router returning a mix of unique and duplicate node names."""
+    import asyncio
+
+    call_count = {"tool_a": 0, "tool_b": 0, "router": 0}
+
+    @node(start=True)
+    async def start(input: dict):
+        return {}
+
+    @router(name="parallel_router")
+    async def parallel_router(input: dict):
+        call_count["router"] += 1
+        if call_count["router"] > 1:
+            return {"choice": "end"}
+        # Mix of unique and duplicate nodes
+        return {"choices": ["tool_a", "tool_b", "tool_a", "tool_b", "tool_a"]}
+
+    @node(name="tool_a")
+    async def tool_a(input: dict):
+        call_count["tool_a"] += 1
+        await asyncio.sleep(0.01)
+        return {"a_count": call_count["tool_a"]}
+
+    @node(name="tool_b")
+    async def tool_b(input: dict):
+        call_count["tool_b"] += 1
+        await asyncio.sleep(0.01)
+        return {"b_count": call_count["tool_b"]}
+
+    @node(end=True)
+    async def end(input: dict):
+        return {"final": "done"}
+
+    g = Graph()
+    g.add_node(start).add_node(parallel_router).add_node(tool_a).add_node(tool_b).add_node(end)
+    g.add_edge(start, parallel_router)
+    g.add_edge(parallel_router, tool_a)
+    g.add_edge(parallel_router, tool_b)
+    g.add_join([tool_a, tool_b], parallel_router)
+    g.add_edge(parallel_router, end)
+
+    result = await g.execute({"input": {}})
+
+    # tool_a should have been called 3 times
+    assert call_count["tool_a"] == 3
+    # tool_b should have been called 2 times
+    assert call_count["tool_b"] == 2
+    # Router called twice
+    assert call_count["router"] == 2
+    assert result["final"] == "done"