Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 16 additions & 11 deletions swarm/graph/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,17 +113,22 @@ async def run(self, inputs: Dict[str, Any],
max_time: int = 600,
return_all_outputs: bool = False) -> List[Any]:

def is_node_useful(node):
if node in self.output_nodes:
return True

for successor in node.successors:
if is_node_useful(successor):
return True
return False

useful_node_ids = [node_id for node_id, node in self.nodes.items() if is_node_useful(node)]
in_degree = {node_id: len(self.nodes[node_id].predecessors) for node_id in useful_node_ids}
# Only ancestors of an output can contribute to the final answer.
# Walk backward once instead of revisiting every forward path.
useful_node_ids = set()
pending = list(self.output_nodes)
while pending:
node = pending.pop()
if node.id in useful_node_ids:
continue
useful_node_ids.add(node.id)
pending.extend(node.predecessors)

in_degree = {
node_id: len(node.predecessors)
for node_id, node in self.nodes.items()
if node_id in useful_node_ids
}
zero_in_degree_queue = [node_id for node_id, deg in in_degree.items() if deg == 0 and node_id in useful_node_ids]

for i, input_node in enumerate(self.input_nodes):
Expand Down
89 changes: 89 additions & 0 deletions test/swarm/graph/test_output_reachability.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
import asyncio
import sys

import pytest

from swarm.graph.composite_graph import CompositeGraph
from swarm.graph.node import Node

pytestmark = pytest.mark.mock_llm


class RecordingNode(Node):
def __init__(self, node_id, executed):
super().__init__("record execution", node_id, combine_inputs_as_one=True)
self.executed = executed

async def _execute(self, inputs, **kwargs):
self.executed.append(self.id)
return {"output": self.id}


def test_shared_ancestors_execute_once_for_multiple_outputs():
executed = []
first = RecordingNode("first", executed)
second = RecordingNode("second", executed)
source = RecordingNode("source", executed)
graph = CompositeGraph(first, "test")
graph.add_node(second)
graph.add_node(source)
graph.output_nodes.append(second)
graph.input_nodes = [source]
source.add_successor(first)
source.add_successor(second)
graph.add_node(RecordingNode("unused", executed))

assert asyncio.run(graph.run({"task": "test"})) == ["first", "second"]
assert executed == ["source", "first", "second"]


def test_deep_graph_does_not_require_recursive_reachability():
executed = []
output = RecordingNode("output", executed)
graph = CompositeGraph(output, "test")
previous = output
for i in range(sys.getrecursionlimit() + 10):
node = graph.add_node(RecordingNode(str(i), executed))
node.add_successor(previous)
previous = node
graph.input_nodes = [previous]

assert asyncio.run(graph.run({"task": "test"})) == ["output"]
assert len(executed) == len(graph.nodes)


def test_dead_diamond_branches_do_not_revisit_every_path():
executed = []
visits = 0

class CountedSuccessors(list):
def __iter__(self):
nonlocal visits
visits += 1
assert visits <= 100, "reachability repeatedly walks the same dead branches"
return super().__iter__()

output = RecordingNode("output", executed)
graph = CompositeGraph(output, "test")
graph.input_nodes = [output]
previous_layer = []
for depth in range(10):
layer = [graph.add_node(RecordingNode(f"{depth}-{i}", executed)) for i in range(2)]
for source in previous_layer:
for target in layer:
source.add_successor(target)
previous_layer = layer
for node in graph.nodes.values():
node.successors = CountedSuccessors(node.successors)

assert asyncio.run(graph.run({"task": "test"})) == ["output"]
assert executed == ["output"]


def test_graph_without_outputs_executes_no_nodes():
executed = []
graph = CompositeGraph(RecordingNode("unused", executed), "test")
graph.output_nodes = []

assert asyncio.run(graph.run({})) == ["No answer since there are no inputs provided"]
assert executed == []