Update aggregator (#20)

andresg4 · web-flow · commit 4b61a7cc198e · 2026-02-18T11:45:00.000-03:00
* update aggregator, gate logic and sequential usage
* updated the diagram
diff --git a/examples/complex_llm_workflow/diagram.md b/examples/complex_llm_workflow/diagram.md
@@ -10,10 +10,10 @@ agent3[Agent3]
 agent4[Agent4]
 in --> agent1
 agent1 --> gate
-gate -->|success| agent2
-gate -->|success| agent3
-gate -->|success| agent4
-gate -->|default| out
+gate -->|failure| out
+gate --> agent2
+gate --> agent3
+gate --> agent4
 agent2 --> parallel_workflow_aggregator
 parallel_workflow_aggregator --> out
 agent3 --> parallel_workflow_aggregator
diff --git a/examples/complex_llm_workflow/generator.rb b/examples/complex_llm_workflow/generator.rb
@@ -85,16 +85,21 @@ def tools
   steps: [llm2, llm3, llm4]
 )
 
+error_workflow = Mars::Workflows::Sequential.new(
+  "Error workflow",
+  steps: []
+)
+
 gate = Mars::Gate.new(
-  condition: ->(input) { input.split.length < 10 ? :success : :error },
+  condition: ->(input) { input.split.length < 10 ? :success : :failure },
   branches: {
-    success: parallel_workflow
+    failure: error_workflow
   }
 )
 
 sequential_workflow = Mars::Workflows::Sequential.new(
   "Sequential workflow",
-  steps: [llm1, gate]
+  steps: [llm1, gate, parallel_workflow]
 )
 
 # Generate and save the diagram
diff --git a/examples/complex_workflow/diagram.md b/examples/complex_workflow/diagram.md
@@ -4,25 +4,23 @@ in((In))
 out((Out))
 agent1[Agent1]
 gate{Gate}
-parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator]
 agent4[Agent4]
 parallel_workflow_aggregator[Parallel workflow Aggregator]
 agent2[Agent2]
 agent3[Agent3]
+parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator]
 agent5[Agent5]
 in --> agent1
 agent1 --> gate
-gate -->|success| agent4
-gate -->|success| agent5
 gate -->|warning| agent4
 gate -->|error| agent2
 gate -->|error| agent3
-gate -->|default| out
+gate --> agent4
+gate --> agent5
 agent4 --> agent2
 agent4 --> agent3
 agent2 --> parallel_workflow_aggregator
 parallel_workflow_aggregator --> parallel_workflow_2_aggregator
-parallel_workflow_aggregator --> out
 agent3 --> parallel_workflow_aggregator
 parallel_workflow_2_aggregator --> out
 agent5 --> parallel_workflow_2_aggregator
diff --git a/examples/complex_workflow/generator.rb b/examples/complex_workflow/generator.rb
@@ -48,7 +48,6 @@ class Agent5 < Mars::Agent
 gate = Mars::Gate.new(
   condition: ->(input) { input[:result] },
   branches: {
-    success: parallel_workflow2,
     warning: sequential_workflow,
     error: parallel_workflow
   }
@@ -57,7 +56,7 @@ class Agent5 < Mars::Agent
 # Create the main workflow: LLM 1 -> Gate
 main_workflow = Mars::Workflows::Sequential.new(
   "Main Pipeline",
-  steps: [llm1, gate]
+  steps: [llm1, gate, parallel_workflow2]
 )
 
 # Generate and save the diagram
diff --git a/lib/mars/aggregator.rb b/lib/mars/aggregator.rb
@@ -8,7 +8,7 @@ def initialize(name = "Aggregator", operation: nil, **kwargs)
       super(**kwargs)
 
       @name = name
-      @operation = operation || ->(inputs) { inputs.join("\n") }
+      @operation = operation || ->(inputs) { inputs }
     end
 
     def run(inputs)
diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb
@@ -15,7 +15,7 @@ def initialize(name = "Gate", condition:, branches:, **kwargs)
     def run(input)
       result = condition.call(input)
 
-      branches[result]&.run(input) || input
+      branches[result] || input
     end
 
     private
diff --git a/lib/mars/rendering/graph/gate.rb b/lib/mars/rendering/graph/gate.rb
@@ -14,8 +14,6 @@ def to_graph(builder, parent_id: nil, value: nil)
             branch.to_graph(builder, parent_id: node_id, value: condition_result)
           end
 
-          builder.add_edge(node_id, "out", "default")
-
           sink_nodes.flatten
         end
       end
diff --git a/lib/mars/rendering/graph/sequential_workflow.rb b/lib/mars/rendering/graph/sequential_workflow.rb
@@ -14,6 +14,8 @@ def to_graph(builder, parent_id: nil, value: nil)
             parent_id = step.node_id
           end
 
+          builder.add_edge(parent_id, "out", value) if sink_nodes.empty?
+
           sink_nodes.flatten
         end
       end
diff --git a/lib/mars/workflows/sequential.rb b/lib/mars/workflows/sequential.rb
@@ -14,7 +14,14 @@ def initialize(name, steps:, **kwargs)
 
       def run(input)
         @steps.each do |step|
-          input = step.run(input)
+          result = step.run(input)
+
+          if result.is_a?(Runnable)
+            input = result.run(input)
+            break
+          else
+            input = result
+          end
         end
 
         input
diff --git a/spec/mars/aggregator_spec.rb b/spec/mars/aggregator_spec.rb
@@ -5,15 +5,9 @@
     context "when called without a block" do
       let(:aggregator) { described_class.new }
 
-      it "joins inputs with newlines" do
-        inputs = %w[first second third]
-        result = aggregator.run(inputs)
-        expect(result).to eq("first\nsecond\nthird")
-      end
-
-      it "handles empty array" do
-        result = aggregator.run([])
-        expect(result).to eq("")
+      it "returns the input as is" do
+        result = aggregator.run([1, 2, 3])
+        expect(result).to eq([1, 2, 3])
       end
     end
 
diff --git a/spec/mars/gate_spec.rb b/spec/mars/gate_spec.rb
@@ -2,29 +2,28 @@
 
 RSpec.describe Mars::Gate do
   describe "#run" do
+    let(:gate) { described_class.new("TestGate", condition: condition, branches: branches) }
+
     context "with simple boolean condition" do
       let(:condition) { ->(input) { input > 5 } }
-      let(:true_branch) { instance_spy(Mars::Runnable) }
       let(:false_branch) { instance_spy(Mars::Runnable) }
-      let(:branches) { { true => true_branch, false => false_branch } }
-      let(:gate) { described_class.new("TestGate", condition: condition, branches: branches) }
-
-      it "executes true branch when condition is true" do
-        allow(true_branch).to receive(:run).with(10).and_return("true result")
+      let(:branches) { { false => false_branch } }
 
+      it "returns the input when no branch matches" do
         result = gate.run(10)
-
-        expect(result).to eq("true result")
-        expect(true_branch).to have_received(:run).with(10)
+        expect(result).to eq(10)
       end
 
-      it "executes false branch when condition is false" do
-        allow(false_branch).to receive(:run).with(3).and_return("false result")
-
+      it "returns the false branch when condition is false" do
         result = gate.run(3)
 
-        expect(result).to eq("false result")
-        expect(false_branch).to have_received(:run).with(3)
+        expect(result).to eq(false_branch)
+      end
+
+      it "does not run the false branch when condition is false" do
+        gate.run(3)
+
+        expect(false_branch).not_to have_received(:run)
       end
     end
 
@@ -33,46 +32,17 @@
       let(:long_branch) { instance_spy(Mars::Runnable) }
       let(:short_branch) { instance_spy(Mars::Runnable) }
       let(:branches) { { "long" => long_branch, "short" => short_branch } }
-      let(:gate) { described_class.new("LengthGate", condition: condition, branches: branches) }
-
-      it "routes to correct branch based on string result" do
-        allow(long_branch).to receive(:run).with("longstring").and_return("long result")
 
+      it "routes to long branch for long strings" do
         result = gate.run("longstring")
 
-        expect(result).to eq("long result")
-        expect(long_branch).to have_received(:run).with("longstring")
+        expect(result).to eq(long_branch)
       end
 
       it "routes to short branch for short strings" do
-        allow(short_branch).to receive(:run).with("hi").and_return("short result")
-
         result = gate.run("hi")
 
-        expect(result).to eq("short result")
-        expect(short_branch).to have_received(:run).with("hi")
-      end
-    end
-
-    context "with missing branch" do
-      let(:condition) { ->(input) { input > 5 ? "high" : "low" } }
-      let(:high_branch) { instance_spy(Mars::Runnable) }
-      let(:branches) { { "high" => high_branch } }
-      let(:gate) { described_class.new("TestGate", condition: condition, branches: branches) }
-
-      it "executes defined branch when condition matches" do
-        allow(high_branch).to receive(:run).with(10).and_return("high result")
-
-        result = gate.run(10)
-
-        expect(result).to eq("high result")
-        expect(high_branch).to have_received(:run).with(10)
-      end
-
-      it "returns the input when branch is not defined" do
-        # For input 3, condition returns "low" which is not in branches
-        result = gate.run(3)
-        expect(result).to eq(3)
+        expect(result).to eq(short_branch)
       end
     end
 
@@ -93,33 +63,21 @@
       let(:branches) { { "low" => low_branch, "medium" => medium_branch, "high" => high_branch } }
 
       it "routes to low branch" do
-        gate = described_class.new("RangeGate", condition: condition, branches: branches)
-        allow(low_branch).to receive(:run).with(5).and_return("low result")
-
         result = gate.run(5)
 
-        expect(result).to eq("low result")
-        expect(low_branch).to have_received(:run).with(5)
+        expect(result).to eq(low_branch)
       end
 
       it "routes to medium branch" do
-        gate = described_class.new("RangeGate", condition: condition, branches: branches)
-        allow(medium_branch).to receive(:run).with(25).and_return("medium result")
-
         result = gate.run(25)
 
-        expect(result).to eq("medium result")
-        expect(medium_branch).to have_received(:run).with(25)
+        expect(result).to eq(medium_branch)
       end
 
       it "routes to high branch" do
-        gate = described_class.new("RangeGate", condition: condition, branches: branches)
-        allow(high_branch).to receive(:run).with(100).and_return("high result")
-
         result = gate.run(100)
 
-        expect(result).to eq("high result")
-        expect(high_branch).to have_received(:run).with(100)
+        expect(result).to eq(high_branch)
       end
     end
   end
diff --git a/spec/mars/workflows/parallel_spec.rb b/spec/mars/workflows/parallel_spec.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 RSpec.describe Mars::Workflows::Parallel do
+  let(:sum_aggregator) { Mars::Aggregator.new("Sum Aggregator", operation: lambda(&:sum)) }
   let(:add_step_class) do
     Class.new do
       def initialize(value)
@@ -45,24 +46,23 @@ def run(_input)
   end
 
   describe "#run" do
-    it "executes steps in parallel" do
+    it "executes steps in parallel without an aggregator" do
       add_five = add_step_class.new(5)
       multiply_three = multiply_step_class.new(3)
       add_two = add_step_class.new(2)
 
       workflow = described_class.new("math_workflow", steps: [add_five, multiply_three, add_two])
 
       # 10 + 5 = 15, 10 * 3 = 30, 10 + 2 = 12
-      expect(workflow.run(10)).to eq("15\n30\n12")
+      expect(workflow.run(10)).to eq([15, 30, 12])
     end
 
     it "executes steps in parallel with a custom aggregator" do
       add_five = add_step_class.new(5)
       multiply_three = multiply_step_class.new(3)
       add_two = add_step_class.new(2)
-      aggregator = Mars::Aggregator.new("Custom Aggregator", operation: lambda(&:sum))
       workflow = described_class.new("math_workflow", steps: [add_five, multiply_three, add_two],
-                                                      aggregator: aggregator)
+                                                      aggregator: sum_aggregator)
 
       expect(workflow.run(10)).to eq(57)
     end
@@ -71,13 +71,13 @@ def run(_input)
       multiply_step = multiply_step_class.new(7)
       workflow = described_class.new("single_step", steps: [multiply_step])
 
-      expect(workflow.run(6)).to eq("42")
+      expect(workflow.run(6)).to eq([42])
     end
 
-    it "returns input unchanged when no steps" do
+    it "returns empty array when no steps" do
       workflow = described_class.new("empty", steps: [])
 
-      expect(workflow.run(42)).to eq("")
+      expect(workflow.run(42)).to eq([])
     end
 
     it "propagates errors from steps" do

Original file line number	Diff line number	Diff line change
`@@ -85,16 +85,21 @@ def tools`
`85`	`85`	`steps: [llm2, llm3, llm4]`
`86`	`86`	`)`
`87`	`87`
	`88`	`+error_workflow = Mars::Workflows::Sequential.new(`
	`89`	`+ "Error workflow",`
	`90`	`+ steps: []`
	`91`	`+)`
	`92`	`+`
`88`	`93`	`gate = Mars::Gate.new(`
`89`		`- condition: ->(input) { input.split.length < 10 ? :success : :error },`
	`94`	`+ condition: ->(input) { input.split.length < 10 ? :success : :failure },`
`90`	`95`	`branches: {`
`91`		`- success: parallel_workflow`
	`96`	`+ failure: error_workflow`
`92`	`97`	`}`
`93`	`98`	`)`
`94`	`99`
`95`	`100`	`sequential_workflow = Mars::Workflows::Sequential.new(`
`96`	`101`	`"Sequential workflow",`
`97`		`- steps: [llm1, gate]`
	`102`	`+ steps: [llm1, gate, parallel_workflow]`
`98`	`103`	`)`
`99`	`104`
`100`	`105`	`# Generate and save the diagram`