Merge pull request #1491 from Steve-Dusty/feat/graph-workflow-on-node-complete

feat: add on_node_complete streaming callback to GraphWorkflow

Author: kyegomez

examples/multi_agent/graphworkflow_examples/graph_workflow_streaming.py

@@ -0,0 +1,130 @@
+"""
+GraphWorkflow Token Streaming Example
+
+Demonstrates real token-by-token streaming from multiple agents.
+You'll see tokens appear in the terminal as each agent generates
+them, with color-coded labels showing which agent is "speaking".
+
+Architecture:
+    Coordinator (Layer 0)
+        -> Market-Analyst   (Layer 1, parallel)
+        -> Tech-Analyst     (Layer 1, parallel)
+        -> Risk-Analyst     (Layer 1, parallel)
+            -> Synthesizer  (Layer 2)
+"""
+
+import sys
+import threading
+import time
+
+from swarms.structs.agent import Agent
+from swarms.structs.graph_workflow import GraphWorkflow
+
+# ANSI colors for each agent
+COLORS = {
+    "Coordinator": "\033[96m",    # cyan
+    "Market-Analyst": "\033[93m", # yellow
+    "Tech-Analyst": "\033[92m",   # green
+    "Risk-Analyst": "\033[91m",   # red
+    "Synthesizer": "\033[95m",    # magenta
+}
+RESET = "\033[0m"
+BOLD = "\033[1m"
+
+# Lock to avoid garbled output from parallel agents
+print_lock = threading.Lock()
+
+
+def create_agent(name: str, description: str) -> Agent:
+    return Agent(
+        agent_name=name,
+        agent_description=description,
+        system_prompt=f"You are {name}. {description} Keep your response to 2-3 sentences.",
+        model_name="gpt-5.4",
+        max_loops=1,
+        verbose=False,
+        print_on=False,
+        streaming_on=True,
+    )
+
+
+def main():
+    # -- Build agents --
+    coordinator = create_agent(
+        "Coordinator",
+        "You coordinate analysis tasks. Briefly outline what each team member should focus on.",
+    )
+    market_analyst = create_agent(
+        "Market-Analyst",
+        "You analyse market trends and competitive landscape.",
+    )
+    tech_analyst = create_agent(
+        "Tech-Analyst",
+        "You evaluate technical feasibility and architecture.",
+    )
+    risk_analyst = create_agent(
+        "Risk-Analyst",
+        "You identify risks and propose mitigations.",
+    )
+    synthesizer = create_agent(
+        "Synthesizer",
+        "You synthesize inputs from multiple analysts into a concise executive summary.",
+    )
+
+    # -- Build workflow --
+    workflow = GraphWorkflow(name="Streaming-Demo")
+    for agent in [coordinator, market_analyst, tech_analyst, risk_analyst, synthesizer]:
+        workflow.add_node(agent)
+
+    workflow.add_edges_from_source(
+        "Coordinator",
+        ["Market-Analyst", "Tech-Analyst", "Risk-Analyst"],
+    )
+    workflow.add_edges_to_target(
+        ["Market-Analyst", "Tech-Analyst", "Risk-Analyst"],
+        "Synthesizer",
+    )
+
+    # -- Token-by-token streaming callback --
+    # Track which agents have printed their header
+    active_agents = {}
+
+    def on_token(node_id: str, token: str) -> None:
+        color = COLORS.get(node_id, "")
+        with print_lock:
+            if node_id not in active_agents:
+                active_agents[node_id] = True
+                sys.stdout.write(f"\n{color}{BOLD}[{node_id}]{RESET}{color} ")
+            sys.stdout.write(f"{color}{token}{RESET}")
+            sys.stdout.flush()
+
+    def on_complete(node_id: str, output) -> None:
+        with print_lock:
+            sys.stdout.write("\n")
+            sys.stdout.flush()
+        # Clear so next run of same agent gets a new header
+        active_agents.pop(node_id, None)
+
+    # -- Run --
+    task = "Evaluate the feasibility of launching an AI-powered personal finance assistant."
+
+    print(f"{BOLD}{'=' * 60}")
+    print("  GraphWorkflow Token Streaming Demo")
+    print(f"{'=' * 60}{RESET}")
+    print(f"\n  Task: {task}\n")
+
+    start = time.time()
+    result = workflow.run(
+        task,
+        streaming_callback=on_token,
+        on_node_complete=on_complete,
+    )
+    elapsed = time.time() - start
+
+    print(f"\n{BOLD}{'=' * 60}")
+    print(f"  Done in {elapsed:.1f}s  |  Agents: {list(result.keys())}")
+    print(f"{'=' * 60}{RESET}")
+
+
+if __name__ == "__main__":
+    main()

swarms/structs/graph_workflow.py

@@ -10,6 +10,7 @@
 from enum import Enum
 from typing import (
     Any,
+    Callable,
     Dict,
     Iterator,
     List,
@@ -741,9 +742,13 @@ def __init__(
         verbose: bool = False,
         backend: str = "networkx",
         checkpoint_dir: Optional[str] = None,
+        on_node_complete: Optional[
+            Callable[[str, Any], None]
+        ] = None,
     ):
         self.id = id
         self.verbose = verbose
+        self.on_node_complete = on_node_complete
 
         if self.verbose:
             logger.info("Initializing GraphWorkflow")
@@ -1713,6 +1718,12 @@ def run(
         self,
         task: Optional[str] = None,
         img: Optional[str] = None,
+        on_node_complete: Optional[
+            Callable[[str, Any], None]
+        ] = None,
+        streaming_callback: Optional[
+            Callable[[str, str], None]
+        ] = None,
         *args: Any,
         **kwargs: Any,
     ) -> Dict[str, Any]:
@@ -1726,6 +1737,14 @@ def run(
         Args:
             task (Optional[str]): Task to execute. Uses self.task if not provided.
             img (Optional[str]): Optional image path for multimodal tasks.
+            on_node_complete (Optional[Callable[[str, Any], None]]): Callback
+                fired immediately when each agent finishes, before the layer
+                completes. Receives ``(node_id, output)``. A callback passed
+                here takes precedence over the instance-level callback set in
+                ``__init__``.
+            streaming_callback (Optional[Callable[[str, str], None]]): Callback
+                fired for every token as agents generate output in real-time.
+                Receives ``(node_id, token)``.
             *args: Additional positional arguments.
             **kwargs: Additional keyword arguments.
 
@@ -1736,6 +1755,11 @@ def run(
                 keyed as ``{node_id}_loop_{loop_number}`` plus the final
                 loop's results under the plain ``node_id`` keys.
         """
+        # Resolve callbacks: run-level overrides instance-level
+        _on_node_complete = (
+            on_node_complete or self.on_node_complete
+        )
+        _streaming_callback = streaming_callback
         run_start_time = time.time()
 
         if task is not None:
@@ -1904,12 +1928,21 @@ def run(
                         # Submit all tasks
                         for node_id, agent, prompt in layer_data:
                             try:
+                                # Build per-agent kwargs, injecting
+                                # streaming_callback if provided.
+                                submit_kwargs = dict(kwargs)
+                                if _streaming_callback is not None:
+                                    _nid = node_id  # capture for closure
+                                    submit_kwargs["streaming_callback"] = (
+                                        lambda token, _nid=_nid: _streaming_callback(_nid, token)
+                                    )
+
                                 future = executor.submit(
                                     agent.run,
                                     prompt,
                                     img,
                                     *args,
-                                    **kwargs,
+                                    **submit_kwargs,
                                 )
                                 future_to_data[future] = (
                                     node_id,
@@ -1982,6 +2015,15 @@ def run(
                                     f"Error adding output to conversation for agent {agent_name}: {e}"
                                 )
 
+                            # Fire the on_node_complete callback
+                            if _on_node_complete is not None:
+                                try:
+                                    _on_node_complete(node_id, output)
+                                except Exception as e:
+                                    logger.exception(
+                                        f"Error in on_node_complete callback for {agent_name}: {e}"
+                                    )
+
                     layer_execution_time = (
                         time.time() - layer_start_time
                     )

tests/structs/test_graph_workflow.py

@@ -931,5 +931,84 @@ def test_graph_workflow_single_loop_backward_compatible():
     ), "Single-loop results should not contain loop-suffixed keys"
 
 
+def test_graph_workflow_on_node_complete_callback_via_run():
+    """Test that on_node_complete callback fires for each agent when passed
+    to run() (fixes #1482)."""
+    agent1 = create_test_agent("Agent1", "Entry agent")
+    agent2 = create_test_agent("Agent2", "End agent")
+
+    workflow = GraphWorkflow(name="Callback-Run-Test")
+    workflow.add_node(agent1)
+    workflow.add_node(agent2)
+    workflow.add_edge(agent1, agent2)
+
+    completed = []
+
+    def on_complete(node_id, output):
+        completed.append((node_id, output))
+
+    result = workflow.run(
+        "Test callback via run",
+        on_node_complete=on_complete,
+    )
+    assert result is not None
+    assert len(completed) == 2
+
+    completed_ids = [nid for nid, _ in completed]
+    assert "Agent1" in completed_ids
+    assert "Agent2" in completed_ids
+
+    # Outputs in callback should match the returned results
+    for node_id, output in completed:
+        assert result[node_id] == output
+
+
+def test_graph_workflow_on_node_complete_callback_via_init():
+    """Test that on_node_complete callback works when set at __init__ level."""
+    agent1 = create_test_agent("Agent1", "Entry agent")
+    agent2 = create_test_agent("Agent2", "End agent")
+
+    completed = []
+
+    def on_complete(node_id, output):
+        completed.append(node_id)
+
+    workflow = GraphWorkflow(
+        name="Callback-Init-Test",
+        on_node_complete=on_complete,
+    )
+    workflow.add_node(agent1)
+    workflow.add_node(agent2)
+    workflow.add_edge(agent1, agent2)
+
+    result = workflow.run("Test callback via init")
+    assert result is not None
+    assert "Agent1" in completed
+    assert "Agent2" in completed
+
+
+def test_graph_workflow_on_node_complete_run_overrides_init():
+    """Test that a callback passed to run() takes precedence over __init__."""
+    agent1 = create_test_agent("Agent1", "Solo agent")
+
+    init_calls = []
+    run_calls = []
+
+    workflow = GraphWorkflow(
+        name="Callback-Override-Test",
+        on_node_complete=lambda nid, out: init_calls.append(nid),
+    )
+    workflow.add_node(agent1)
+
+    workflow.run(
+        "Test override",
+        on_node_complete=lambda nid, out: run_calls.append(nid),
+    )
+
+    # Only the run-level callback should have been called
+    assert len(run_calls) == 1
+    assert len(init_calls) == 0
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])