feat: add join edge and release notes

Author: jamescalam

docs/user-guide/release-notes.md

@@ -4,6 +4,16 @@ See below for all notable changes to the GraphAI library.
 
 ### Added
 - Direct Starlette support for `GraphEvent` and `EventCallback` objects
+- Parallel node execution for graphs with multiple outgoing edges from a single node
+  - Automatic concurrent execution when a node has multiple successors
+  - State merging from parallel branches
+  - Configurable through standard `add_edge()` calls
+- New `add_join()` method for explicit convergence of parallel branches
+  - Synchronizes multiple parallel branches to a single destination node
+  - Ensures convergence node executes only once with merged state
+  - Prevents duplicate execution of downstream nodes
+- New `add_parallel()` convenience method for creating parallel branches
+  - Syntactic sugar for adding multiple edges from one source to multiple destinations
 
 ## [0.0.9] - 2025-09-05
 

graphai/graph.py

@@ -2,6 +2,7 @@
 import asyncio
 from typing import Any, Iterable, Protocol
 from graphlib import TopologicalSorter, CycleError
+
 from graphai.callback import Callback
 from graphai.utils import logger
 
@@ -64,6 +65,7 @@ def __init__(
         self.edges: list[Any] = []
         self.start_node: NodeProtocol | None = None
         self.end_nodes: list[NodeProtocol] = []
+        self.join_nodes: set[NodeProtocol] = set()
         self.Callback: type[Callback] = Callback
         self.max_steps = max_steps
         self.state = initial_state or {}
@@ -130,6 +132,18 @@ def add_node(self, node: NodeProtocol) -> Graph:
             self.end_nodes.append(node)
         return self
 
+    def _get_node(self, node_candidate: NodeProtocol | str) -> NodeProtocol:
+        # first get node from graph
+        if isinstance(node_candidate, str):
+            node = self.nodes.get(node_candidate)
+        else:
+            # check if it's a node-like object by looking for required attributes
+            if hasattr(node_candidate, "name"):
+                node = self.nodes.get(node_candidate.name)
+        if node is None:
+            raise ValueError(f"Node with name '{node_candidate}' not found.")
+        return node
+
     def add_edge(
         self, source: NodeProtocol | str, destination: NodeProtocol | str
     ) -> Graph:
@@ -141,33 +155,10 @@ def add_edge(
         """
         source_node, destination_node = None, None
         # get source node from graph
-        source_name: str
-        if isinstance(source, str):
-            source_node = self.nodes.get(source)
-            source_name = source
-        else:
-            # Check if it's a node-like object by looking for required attributes
-            if hasattr(source, "name"):
-                source_node = self.nodes.get(source.name)
-                source_name = source.name
-            else:
-                source_name = str(source)
-        if source_node is None:
-            raise ValueError(f"Node with name '{source_name}' not found.")
+        source_node = self._get_node(node_candidate=source)
         # get destination node from graph
-        destination_name: str
-        if isinstance(destination, str):
-            destination_node = self.nodes.get(destination)
-            destination_name = destination
-        else:
-            # Check if it's a node-like object by looking for required attributes
-            if hasattr(destination, "name"):
-                destination_node = self.nodes.get(destination.name)
-                destination_name = destination.name
-            else:
-                destination_name = str(destination)
-        if destination_node is None:
-            raise ValueError(f"Node with name '{destination_name}' not found.")
+        destination_node = self._get_node(node_candidate=destination)
+        # create edge
         edge = Edge(source_node, destination_node)
         self.edges.append(edge)
         return self
@@ -214,7 +205,6 @@ def compile(self, *, strict: bool = False) -> Graph:
         nodes = getattr(self, "nodes", None)
         if not isinstance(nodes, dict) or not nodes:
             raise GraphCompileError("No nodes have been added to the graph")
-
         start_name: str | None = None
         # Bind and narrow the attribute for mypy
         start_node: _HasName | None = getattr(self, "start_node", None)
@@ -230,21 +220,17 @@ def compile(self, *, strict: bool = False) -> Graph:
                 raise GraphCompileError(f"Multiple start nodes defined: {starts}")
             if len(starts) == 1:
                 start_name = starts[0]
-
         if not start_name:
             raise GraphCompileError("No start node defined")
-
         # at least one end node
         if not any(
             getattr(n, "is_end", False) or getattr(n, "end", False)
             for n in nodes.values()
         ):
             raise GraphCompileError("No end node defined")
-
         # normalize edges into adjacency {src: set(dst)}
         raw_edges = getattr(self, "edges", None)
         adj: dict[str, set[str]] = {name: set() for name in nodes.keys()}
-
         def _add_edge(src: str, dst: str) -> None:
             if src not in nodes:
                 raise GraphCompileError(f"Edge references unknown source node: {src}")
@@ -253,7 +239,6 @@ def _add_edge(src: str, dst: str) -> None:
                     f"Edge from {src} references unknown node(s): ['{dst}']"
                 )
             adj[src].add(dst)
-
         if raw_edges is None:
             pass
         elif isinstance(raw_edges, dict):
@@ -273,13 +258,11 @@ def _add_edge(src: str, dst: str) -> None:
                 iterator = iter(raw_edges)
             except TypeError:
                 raise GraphCompileError("Internal edge map has unsupported type")
-
             for item in iterator:
                 # (src, dst) OR (src, Iterable[dst])
                 if isinstance(item, (tuple, list)) and len(item) == 2:
                     raw_src, rhs = item
                     src = _require_name(raw_src, "source")
-
                     if isinstance(rhs, str) or getattr(rhs, "name", None):
                         dst = _require_name(rhs, "destination")
                         _add_edge(src, rhs)
@@ -294,7 +277,6 @@ def _add_edge(src: str, dst: str) -> None:
                                 "Edge tuple second item must be a destination or an iterable of destinations"
                             )
                     continue
-
                 # Mapping-style: {"source": "...", "destination": "..."} or {"src": "...", "dst": "..."}
                 if isinstance(item, dict):
                     src = _require_name(item.get("source", item.get("src")), "source")
@@ -303,7 +285,6 @@ def _add_edge(src: str, dst: str) -> None:
                     )
                     _add_edge(src, dst)
                     continue
-
                 # Object with attributes .source/.destination (or .src/.dst)
                 if hasattr(item, "source") or hasattr(item, "src"):
                     src = _require_name(
@@ -315,13 +296,11 @@ def _add_edge(src: str, dst: str) -> None:
                     )
                     _add_edge(src, dst)
                     continue
-
                 # If none matched, this is an unsupported edge record
                 raise GraphCompileError(
                     "Edges must be dict[str, Iterable[str]] or an iterable of (src, dst), "
                     "(src, Iterable[dst]), mapping{'source'/'destination'}, or objects with .source/.destination"
                 )
-
         # reachability from start
         seen: set[str] = set()
         stack = [start_name]
@@ -331,11 +310,9 @@ def _add_edge(src: str, dst: str) -> None:
                 continue
             seen.add(cur)
             stack.extend(adj.get(cur, ()))
-
         unreachable = sorted(set(nodes.keys()) - seen)
         if unreachable:
             raise GraphCompileError(f"Unreachable nodes: {unreachable}")
-
         # optional cycle detection (strict mode)
         if strict:
             preds: dict[str, set[str]] = {n: set() for n in nodes.keys()}
@@ -346,7 +323,6 @@ def _add_edge(src: str, dst: str) -> None:
                 list(TopologicalSorter(preds).static_order())
             except CycleError as e:
                 raise GraphCompileError("cycle detected in graph (strict mode)") from e
-
         return self
 
     def _validate_output(self, output: dict[str, Any], node_name: str):
@@ -358,7 +334,13 @@ def _validate_output(self, output: dict[str, Any], node_name: str):
 
     def _get_next_nodes(self, current_node: NodeProtocol) -> list[NodeProtocol]:
         """Return all successor nodes for the given node."""
-        return [edge.destination for edge in self.edges if edge.source == current_node]
+        # we skip JoinEdge because they don't have regular destinations
+        # and next nodes for those are handled in the execute method
+        return [
+            edge.destination
+            for edge in self.edges
+            if isinstance(edge, Edge) and edge.source == current_node
+        ]
 
     async def _invoke_node(
         self, node: NodeProtocol, state: dict[str, Any], callback: Callback
@@ -379,6 +361,7 @@ async def _execute_branch(
         state: dict[str, Any],
         callback: Callback,
         steps: int,
+        stop_at_join: bool = False,
     ):
         """Recursively execute a branch starting from `current_node`.
         When a node has multiple successors, run them concurrently and merge their outputs."""
@@ -392,6 +375,9 @@ async def _execute_branch(
                 del output["choice"]
                 current_node = self._get_node_by_name(node_name=next_node_name)
                 continue
+            if stop_at_join and current_node in self.join_nodes:
+                # for parallel branches, wait at JoinEdge until all branches are complete
+                return state
 
             next_nodes = self._get_next_nodes(current_node)
             if not next_nodes:
@@ -404,17 +390,43 @@ async def _execute_branch(
                 # Run each branch concurrently
                 results = await asyncio.gather(
                     *[
-                        self._execute_branch(n, state.copy(), callback, steps + 1)
+                        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
+                        )
                         for n in next_nodes
                     ]
                 )
+                # merge states returned by each branch
                 merged = state.copy()
                 for res in results:
-                    # merge states returned by each branch
                     for k, v in res.items():
                         if k != "callback":
                             merged[k] = v
-                return merged
+                if set(next_nodes) & self.join_nodes:
+                    # if any of the next nodes are join nodes, we need to continue from the
+                    # JoinEdge.destination node
+                    join_edge = next(
+                        (
+                            e for e in self.edges if isinstance(e, JoinEdge)
+                            and any(n in e.sources for n in next_nodes)
+                        ),
+                        None
+                    )
+                    if not join_edge:
+                        raise Exception("No JoinEdge found for next_nodes")
+                    # set current_node (for next iteration) to the JoinEdge.destination
+                    current_node = join_edge.destination
+                    # continue to the destination node with our merged state
+                    state = merged
+                    continue
+                else:
+                    # if this happens we have multiple branches that do not join so we
+                    # can just return the merged states
+                    return merged
             steps += 1
             if steps >= self.max_steps:
                 raise Exception(
@@ -502,20 +514,46 @@ def _get_node_by_name(self, node_name: str) -> NodeProtocol:
 
     def _get_next_node(self, current_node):
         for edge in self.edges:
-            if edge.source == current_node:
+            if isinstance(edge, Edge) and edge.source == current_node:
                 return edge.destination
+            # we skip JoinEdge because they don't have regular destinations
+            # and next nodes for those are handled in the execute method
         raise Exception(
             f"No outgoing edge found for current node '{current_node.name}'."
         )
 
     def add_parallel(
         self, source: NodeProtocol | str, destinations: list[NodeProtocol | str]
     ):
-        """Add multiple outgoing edges from a single source node to be executed in parallel."""
+        """Add multiple outgoing edges from a single source node to be executed in parallel.
+
+        Args:
+            source: The source node for the parallel branches.
+            destinations: The list of destination nodes for the parallel branches.
+        """
         for dest in destinations:
             self.add_edge(source, dest)
         return self
 
+    def add_join(
+        self, sources: list[NodeProtocol | str], destination: NodeProtocol | str
+    ):
+        """Joins multiple parallel branches into a single branch.
+        
+        Args:
+            sources: The list of source nodes for the join.
+            destination: The destination node for the join.
+        """
+        # get source nodes from graph
+        source_nodes = [self._get_node(node_candidate=source) for source in sources]
+        # get destination node from graph
+        destination_node = self._get_node(node_candidate=destination)
+        # create join edge
+        edge = JoinEdge(source_nodes, destination_node)
+        self.edges.append(edge)
+        self.join_nodes.update(source_nodes)
+        return self
+
     def visualize(self, *, save_path: str | None = None):
         """Render the current graph. If matplotlib is not installed,
         raise a helpful error telling users to install the viz extra.
@@ -611,3 +649,8 @@ class Edge:
     def __init__(self, source, destination):
         self.source = source
         self.destination = destination
+
+class JoinEdge:
+    def __init__(self, sources, destination):
+        self.sources = sources
+        self.destination = destination