Skip to content

Commit a5c12d5

Browse files
santibsantib
committed
Implement the rest of the plan
1 parent ed501ad commit a5c12d5

10 files changed

Lines changed: 403 additions & 70 deletions

File tree

‎lib/mars/execution_context.rb‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,11 @@ def initialize(input: nil, global_state: {})
1111
end
1212

1313
def [](step_name)
14-
outputs[step_name]
14+
outputs[step_name.to_sym]
1515
end
1616

1717
def record(step_name, output)
18-
@outputs[step_name] = output
18+
@outputs[step_name.to_sym] = output
1919
@current_input = output
2020
end
2121

‎lib/mars/rendering/graph.rb‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ module MARS
44
module Rendering
55
module Graph
66
def self.include_extensions
7+
MARS::Runnable.include(Runnable)
78
MARS::AgentStep.include(AgentStep)
89
MARS::Gate.include(Gate)
910
MARS::Workflows::Sequential.include(SequentialWorkflow)

‎lib/mars/rendering/graph/runnable.rb‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# frozen_string_literal: true
2+
3+
module MARS
4+
module Rendering
5+
module Graph
6+
module Runnable
7+
include Base
8+
9+
def to_graph(builder, parent_id: nil, value: nil)
10+
builder.add_node(node_id, name, Node::STEP)
11+
builder.add_edge(parent_id, node_id, value)
12+
13+
[node_id]
14+
end
15+
end
16+
end
17+
end
18+
end

‎lib/mars/rendering/html.rb‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# frozen_string_literal: true
2+
3+
module MARS
4+
module Rendering
5+
class Html
6+
MERMAID_CDN = "https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"
7+
8+
def initialize(obj)
9+
@mermaid = Mermaid.new(obj)
10+
end
11+
12+
def render
13+
diagram = @mermaid.graph_mermaid.join("\n")
14+
direction = "LR"
15+
16+
<<~HTML
17+
<!DOCTYPE html>
18+
<html lang="en">
19+
<head>
20+
<meta charset="UTF-8">
21+
<meta name="viewport" content="width=device-width, initial-scale=1.0">
22+
<title>MARS Workflow</title>
23+
<script src="#{MERMAID_CDN}"></script>
24+
</head>
25+
<body>
26+
<pre class="mermaid">
27+
flowchart #{direction}
28+
#{diagram}
29+
</pre>
30+
<script>mermaid.initialize({ startOnLoad: true });</script>
31+
</body>
32+
</html>
33+
HTML
34+
end
35+
end
36+
end
37+
end

‎lib/mars/workflows/parallel.rb‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,15 @@ def initialize(name, steps:, aggregator: nil, **kwargs)
1111
end
1212

1313
def run(input)
14+
context = ensure_context(input)
1415
errors = []
15-
results = execute_steps(input, errors)
16+
child_contexts = []
17+
results = execute_steps(context, errors, child_contexts)
1618

1719
raise AggregateError, errors if errors.any?
1820

21+
context.merge(child_contexts)
22+
1923
has_global_halt = results.any? { |r| r.is_a?(Halt) && r.global? }
2024
unwrapped = results.map { |r| r.is_a?(Halt) ? r.result : r }
2125
result = aggregator.run(unwrapped)
@@ -26,11 +30,27 @@ def run(input)
2630

2731
attr_reader :steps, :aggregator
2832

29-
def execute_steps(input, errors)
33+
def execute_steps(context, errors, child_contexts)
3034
Async do |workflow|
3135
tasks = steps.map do |step|
36+
child_ctx = context.fork
37+
child_contexts << child_ctx
38+
3239
workflow.async do
33-
step.run(input)
40+
step.run_before_hooks(child_ctx)
41+
42+
step_input = step.formatter.format_input(child_ctx)
43+
result = step.run(step_input)
44+
45+
if result.is_a?(Halt)
46+
step.run_after_hooks(child_ctx, result)
47+
result
48+
else
49+
formatted = step.formatter.format_output(result)
50+
child_ctx.record(step.name, formatted)
51+
step.run_after_hooks(child_ctx, formatted)
52+
formatted
53+
end
3454
rescue StandardError => e
3555
errors << { error: e, step_name: step.name }
3656
end
@@ -39,6 +59,10 @@ def execute_steps(input, errors)
3959
tasks.map(&:wait)
4060
end.result
4161
end
62+
63+
def ensure_context(input)
64+
input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
65+
end
4266
end
4367
end
4468
end

‎lib/mars/workflows/sequential.rb‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,23 +10,41 @@ def initialize(name, steps:, **kwargs)
1010
end
1111

1212
def run(input)
13-
@steps.each do |step|
14-
input = step.run(input)
15-
16-
next unless input.is_a?(Halt)
17-
18-
return input if input.global?
13+
context = ensure_context(input)
1914

20-
input = input.result
21-
break
15+
@steps.each do |step|
16+
step.run_before_hooks(context)
17+
18+
step_input = step.formatter.format_input(context)
19+
result = step.run(step_input)
20+
21+
if result.is_a?(Halt)
22+
if result.global?
23+
step.run_after_hooks(context, result)
24+
return result
25+
end
26+
27+
formatted = step.formatter.format_output(result.result)
28+
context.record(step.name, formatted)
29+
step.run_after_hooks(context, formatted)
30+
break
31+
end
32+
33+
formatted = step.formatter.format_output(result)
34+
context.record(step.name, formatted)
35+
step.run_after_hooks(context, formatted)
2236
end
2337

24-
input
38+
context.current_input
2539
end
2640

2741
private
2842

2943
attr_reader :steps
44+
45+
def ensure_context(input)
46+
input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
47+
end
3048
end
3149
end
3250
end

‎spec/mars/rendering/html_spec.rb‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# frozen_string_literal: true
2+
3+
RSpec.describe MARS::Rendering::Html do
4+
it "renders a self-contained HTML page with mermaid diagram" do
5+
step1 = MARS::AgentStep.new(name: "step1")
6+
step2 = MARS::AgentStep.new(name: "step2")
7+
workflow = MARS::Workflows::Sequential.new("pipeline", steps: [step1, step2])
8+
9+
html = described_class.new(workflow).render
10+
11+
expect(html).to include("<!DOCTYPE html>")
12+
expect(html).to include("mermaid")
13+
expect(html).to include("flowchart LR")
14+
expect(html).to include("step1")
15+
expect(html).to include("step2")
16+
expect(html).to include("mermaid.initialize")
17+
end
18+
19+
it "includes the Mermaid CDN script" do
20+
step = MARS::AgentStep.new(name: "step")
21+
workflow = MARS::Workflows::Sequential.new("simple", steps: [step])
22+
23+
html = described_class.new(workflow).render
24+
25+
expect(html).to include(MARS::Rendering::Html::MERMAID_CDN)
26+
end
27+
end

‎spec/mars/rendering/mermaid_spec.rb‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# frozen_string_literal: true
2+
3+
RSpec.describe MARS::Rendering::Mermaid do
4+
it "renders a custom Runnable subclass as a box node" do
5+
step = Class.new(MARS::Runnable) do
6+
def run(input) = input
7+
end.new(name: "custom_step")
8+
9+
mermaid = described_class.new(step)
10+
output = mermaid.render
11+
12+
expect(output).to include("custom_step[custom_step]")
13+
end
14+
15+
it "renders an AgentStep as a box node" do
16+
step = MARS::AgentStep.new(name: "my_agent")
17+
mermaid = described_class.new(step)
18+
output = mermaid.render
19+
20+
expect(output).to include("my_agent[my_agent]")
21+
end
22+
23+
it "renders a Gate as a diamond node" do
24+
gate = MARS::Gate.new(
25+
"my_gate",
26+
check: ->(_input) { :branch },
27+
fallbacks: {
28+
branch: Class.new(MARS::Runnable) do
29+
def run(input) = input
30+
end.new(name: "branch_step")
31+
}
32+
)
33+
34+
mermaid = described_class.new(gate)
35+
output = mermaid.render
36+
37+
expect(output).to include("my_gate{my_gate}")
38+
expect(output).to include("|branch|")
39+
end
40+
41+
it "renders a Sequential workflow with subgraph" do
42+
step1 = MARS::AgentStep.new(name: "step1")
43+
step2 = MARS::AgentStep.new(name: "step2")
44+
workflow = MARS::Workflows::Sequential.new("pipeline", steps: [step1, step2])
45+
46+
mermaid = described_class.new(workflow)
47+
output = mermaid.render
48+
49+
expect(output).to include("subgraph pipeline")
50+
expect(output).to include("step1[step1]")
51+
expect(output).to include("step2[step2]")
52+
end
53+
54+
it "renders a Parallel workflow with aggregator" do
55+
step1 = MARS::AgentStep.new(name: "step1")
56+
step2 = MARS::AgentStep.new(name: "step2")
57+
workflow = MARS::Workflows::Parallel.new("parallel", steps: [step1, step2])
58+
59+
mermaid = described_class.new(workflow)
60+
output = mermaid.render
61+
62+
expect(output).to include("subgraph parallel")
63+
expect(output).to include("step1[step1]")
64+
expect(output).to include("step2[step2]")
65+
expect(output).to include("parallel_aggregator")
66+
end
67+
end

0 commit comments

Comments
 (0)