From 60d38abba538892d7d32676d796df9ed1ea647b6 Mon Sep 17 00:00:00 2001 From: Ignacio Perez Date: Thu, 16 Oct 2025 18:22:28 -0300 Subject: [PATCH] Change rendering approach --- examples/complex_workflow/diagram.md | 29 ++++++++++ examples/complex_workflow/generator.rb | 55 +++++++++++++++++++ examples/parallel_workflow/diagram.md | 16 ++++++ examples/parallel_workflow/generator.rb | 2 +- examples/simple_workflow/diagram.md | 21 ++++--- examples/simple_workflow/generator.rb | 2 +- lib/mars.rb | 2 +- lib/mars/exit.rb | 15 ----- lib/mars/gate.rb | 2 +- lib/mars/rendering/graph.rb | 15 +++++ lib/mars/rendering/graph/agent.rb | 18 ++++++ lib/mars/rendering/graph/aggregator.rb | 17 ++++++ lib/mars/rendering/graph/base.rb | 32 +++++++++++ lib/mars/rendering/graph/builder.rb | 30 ++++++++++ lib/mars/rendering/graph/gate.rb | 24 ++++++++ lib/mars/rendering/graph/node.rb | 22 ++++++++ lib/mars/rendering/graph/parallel_workflow.rb | 24 ++++++++ .../rendering/graph/sequential_workflow.rb | 22 ++++++++ lib/mars/rendering/mermaid.rb | 53 ++++++++++++++---- lib/mars/rendering/mermaid/agent.rb | 19 ------- lib/mars/rendering/mermaid/aggregator.rb | 20 ------- lib/mars/rendering/mermaid/base.rb | 17 ------ lib/mars/rendering/mermaid/exit.rb | 15 ----- lib/mars/rendering/mermaid/gate.rb | 28 ---------- .../rendering/mermaid/parallel_workflow.rb | 16 ------ .../rendering/mermaid/sequential_workflow.rb | 41 -------------- lib/mars/workflows/parallel.rb | 2 +- 27 files changed, 363 insertions(+), 196 deletions(-) create mode 100644 examples/complex_workflow/diagram.md create mode 100755 examples/complex_workflow/generator.rb create mode 100644 examples/parallel_workflow/diagram.md delete mode 100644 lib/mars/exit.rb create mode 100644 lib/mars/rendering/graph.rb create mode 100644 lib/mars/rendering/graph/agent.rb create mode 100644 lib/mars/rendering/graph/aggregator.rb create mode 100644 lib/mars/rendering/graph/base.rb create mode 100644 lib/mars/rendering/graph/builder.rb create mode 100644 lib/mars/rendering/graph/gate.rb create mode 100644 lib/mars/rendering/graph/node.rb create mode 100644 lib/mars/rendering/graph/parallel_workflow.rb create mode 100644 lib/mars/rendering/graph/sequential_workflow.rb delete mode 100644 lib/mars/rendering/mermaid/agent.rb delete mode 100644 lib/mars/rendering/mermaid/aggregator.rb delete mode 100644 lib/mars/rendering/mermaid/base.rb delete mode 100644 lib/mars/rendering/mermaid/exit.rb delete mode 100644 lib/mars/rendering/mermaid/gate.rb delete mode 100644 lib/mars/rendering/mermaid/parallel_workflow.rb delete mode 100644 lib/mars/rendering/mermaid/sequential_workflow.rb diff --git a/examples/complex_workflow/diagram.md b/examples/complex_workflow/diagram.md new file mode 100644 index 0000000..32e44d7 --- /dev/null +++ b/examples/complex_workflow/diagram.md @@ -0,0 +1,29 @@ +```mermaid +flowchart LR +in((In)) +out((Out)) +llm_1[LLM 1] +gate{Gate} +parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator] +llm_4[LLM 4] +parallel_workflow_aggregator[Parallel workflow Aggregator] +llm_2[LLM 2] +llm_3[LLM 3] +llm_5[LLM 5] +in --> llm_1 +llm_1 --> gate +gate -->|success| llm_4 +gate -->|success| llm_5 +gate -->|warning| llm_4 +gate -->|error| llm_2 +gate -->|error| llm_3 +gate -->|default| out +llm_4 --> llm_2 +llm_4 --> llm_3 +llm_2 --> parallel_workflow_aggregator +parallel_workflow_aggregator --> parallel_workflow_2_aggregator +parallel_workflow_aggregator --> out +llm_3 --> parallel_workflow_aggregator +parallel_workflow_2_aggregator --> out +llm_5 --> parallel_workflow_2_aggregator +``` diff --git a/examples/complex_workflow/generator.rb b/examples/complex_workflow/generator.rb new file mode 100755 index 0000000..1acdd55 --- /dev/null +++ b/examples/complex_workflow/generator.rb @@ -0,0 +1,55 @@ +#!/usr/bin/env ruby +# frozen_string_literal: true + +require_relative "../../lib/mars" + +# Create the LLMs +llm1 = Mars::Agent.new(name: "LLM 1") + +llm2 = Mars::Agent.new(name: "LLM 2") + +llm3 = Mars::Agent.new(name: "LLM 3") + +llm4 = Mars::Agent.new(name: "LLM 4") + +llm5 = Mars::Agent.new(name: "LLM 5") + +# Create a parallel workflow (LLM 2 x LLM 3) +parallel_workflow = Mars::Workflows::Parallel.new( + "Parallel workflow", + steps: [llm2, llm3] +) + +# Create a sequential workflow (Parallel workflow -> LLM 4) +sequential_workflow = Mars::Workflows::Sequential.new( + "Sequential workflow", + steps: [llm4, parallel_workflow] +) + +# Create a parallel workflow (Sequential workflow x LLM 5) +parallel_workflow2 = Mars::Workflows::Parallel.new( + "Parallel workflow 2", + steps: [sequential_workflow, llm5] +) + +# Create the gate that decides between exit or continue +gate = Mars::Gate.new( + name: "Gate", + condition: ->(input) { input[:result] }, + branches: { + success: parallel_workflow2, + warning: sequential_workflow, + error: parallel_workflow + } +) + +# Create the main workflow: LLM 1 -> Gate +main_workflow = Mars::Workflows::Sequential.new( + "Main Pipeline", + steps: [llm1, gate] +) + +# Generate and save the diagram +diagram = Mars::Rendering::Mermaid.new(main_workflow).render +File.write("examples/complex_workflow/diagram.md", diagram) +puts "Complex workflow diagram saved to: examples/complex_workflow/diagram.md" diff --git a/examples/parallel_workflow/diagram.md b/examples/parallel_workflow/diagram.md new file mode 100644 index 0000000..24a4829 --- /dev/null +++ b/examples/parallel_workflow/diagram.md @@ -0,0 +1,16 @@ +```mermaid +flowchart LR +in((In)) +out((Out)) +parallel_workflow_aggregator[Parallel workflow Aggregator] +llm_1[LLM 1] +llm_2[LLM 2] +llm_3[LLM 3] +in --> llm_1 +in --> llm_2 +in --> llm_3 +llm_1 --> parallel_workflow_aggregator +parallel_workflow_aggregator --> out +llm_2 --> parallel_workflow_aggregator +llm_3 --> parallel_workflow_aggregator +``` diff --git a/examples/parallel_workflow/generator.rb b/examples/parallel_workflow/generator.rb index 8873688..069b542 100755 --- a/examples/parallel_workflow/generator.rb +++ b/examples/parallel_workflow/generator.rb @@ -17,6 +17,6 @@ ) # Generate and save the diagram -diagram = Mars::Rendering::Mermaid.render(parallel_workflow) +diagram = Mars::Rendering::Mermaid.new(parallel_workflow).render File.write("examples/parallel_workflow/diagram.md", diagram) puts "Parallel workflow diagram saved to: examples/parallel_workflow/diagram.md" diff --git a/examples/simple_workflow/diagram.md b/examples/simple_workflow/diagram.md index b7f90b2..a2ccca5 100644 --- a/examples/simple_workflow/diagram.md +++ b/examples/simple_workflow/diagram.md @@ -1,12 +1,15 @@ ```mermaid flowchart LR -In(("In")) --> -LLM_1["LLM 1"] -LLM_1 --> Gate -Gate{"Gate"} -Gate -->|success| LLM_2["LLM 2"] -LLM_2 --> LLM_3 -LLM_3["LLM 3"] -LLM_3 --> Out(("Out")) -Gate -->|default| exit((Exit)) +in((In)) +out((Out)) +llm_1[LLM 1] +gate{Gate} +llm_2[LLM 2] +llm_3[LLM 3] +in --> llm_1 +llm_1 --> gate +gate -->|success| llm_2 +gate -->|default| out +llm_2 --> llm_3 +llm_3 --> out ``` diff --git a/examples/simple_workflow/generator.rb b/examples/simple_workflow/generator.rb index 31841df..f67a681 100755 --- a/examples/simple_workflow/generator.rb +++ b/examples/simple_workflow/generator.rb @@ -32,6 +32,6 @@ ) # Generate and save the diagram -diagram = Mars::Rendering::Mermaid.render(main_workflow) +diagram = Mars::Rendering::Mermaid.new(main_workflow).render File.write("examples/simple_workflow/diagram.md", diagram) puts "Simple workflow diagram saved to: examples/simple_workflow/diagram.md" diff --git a/lib/mars.rb b/lib/mars.rb index 1fd23b4..831d67c 100644 --- a/lib/mars.rb +++ b/lib/mars.rb @@ -9,4 +9,4 @@ module Mars class Error < StandardError; end end -Mars::Rendering::Mermaid.include_extensions +Mars::Rendering::Graph.include_extensions diff --git a/lib/mars/exit.rb b/lib/mars/exit.rb deleted file mode 100644 index 9ab9b23..0000000 --- a/lib/mars/exit.rb +++ /dev/null @@ -1,15 +0,0 @@ -# frozen_string_literal: true - -module Mars - class Exit < Runnable - attr_reader :name - - def initialize(name: "Exit") - @name = name - end - - def run(input) - input - end - end -end diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb index 1223b49..f838912 100644 --- a/lib/mars/gate.rb +++ b/lib/mars/gate.rb @@ -7,7 +7,7 @@ class Gate < Runnable def initialize(name:, condition:, branches:) @name = name @condition = condition - @branches = Hash.new(Exit.new).merge(branches) + @branches = branches end def run(input) diff --git a/lib/mars/rendering/graph.rb b/lib/mars/rendering/graph.rb new file mode 100644 index 0000000..916d9dd --- /dev/null +++ b/lib/mars/rendering/graph.rb @@ -0,0 +1,15 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + def self.include_extensions + Mars::Agent.include(Agent) + Mars::Gate.include(Gate) + Mars::Workflows::Sequential.include(SequentialWorkflow) + Mars::Workflows::Parallel.include(ParallelWorkflow) + Mars::Aggregator.include(Aggregator) + end + end + end +end diff --git a/lib/mars/rendering/graph/agent.rb b/lib/mars/rendering/graph/agent.rb new file mode 100644 index 0000000..e4fd77a --- /dev/null +++ b/lib/mars/rendering/graph/agent.rb @@ -0,0 +1,18 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + module Agent + include Base + + def to_graph(builder, parent_id: nil, value: nil) + builder.add_node(node_id, name, Node::STEP) + builder.add_edge(parent_id, node_id, value) + + [node_id] + end + end + end + end +end diff --git a/lib/mars/rendering/graph/aggregator.rb b/lib/mars/rendering/graph/aggregator.rb new file mode 100644 index 0000000..e134f44 --- /dev/null +++ b/lib/mars/rendering/graph/aggregator.rb @@ -0,0 +1,17 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + module Aggregator + include Base + + def to_graph(builder, parent_id: nil, value: nil) + builder.add_edge(parent_id, node_id, value) + + [node_id] + end + end + end + end +end diff --git a/lib/mars/rendering/graph/base.rb b/lib/mars/rendering/graph/base.rb new file mode 100644 index 0000000..7f8068c --- /dev/null +++ b/lib/mars/rendering/graph/base.rb @@ -0,0 +1,32 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + module Base + def build_graph(builder = Mars::Rendering::Graph::Builder.new) + builder.add_node("in", "In", Node::INPUT) + builder.add_node("out", "Out", Node::OUTPUT) + + sink_nodes = to_graph(builder, parent_id: "in") + + sink_nodes.each do |sink_node| + builder.add_edge(sink_node, "out") + end + + [builder.adjacency, builder.nodes] + end + + def node_id + @node_id ||= sanitize(name) + end + + private + + def sanitize(name) + name.to_s.gsub(/[^a-zA-Z0-9]/, "_").downcase + end + end + end + end +end diff --git a/lib/mars/rendering/graph/builder.rb b/lib/mars/rendering/graph/builder.rb new file mode 100644 index 0000000..36d84e0 --- /dev/null +++ b/lib/mars/rendering/graph/builder.rb @@ -0,0 +1,30 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + class Builder + attr_reader :adjacency, :nodes + + def initialize + @adjacency = Hash.new { |h, k| h[k] = [] } + @nodes = {} + end + + def add_edge(from, to, value = nil) + return unless from && to + + # can we avoid visiting the node twice instead? + adjacency[from] << [to, value] unless adjacency[from].include?([to, value]) + adjacency[to] = [] unless adjacency[to] + end + + def add_node(id, value, type) + return if nodes.key?(id) + + nodes[id] = Node.new(id, value, type) + end + end + end + end +end diff --git a/lib/mars/rendering/graph/gate.rb b/lib/mars/rendering/graph/gate.rb new file mode 100644 index 0000000..f35f06a --- /dev/null +++ b/lib/mars/rendering/graph/gate.rb @@ -0,0 +1,24 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + module Gate + include Base + + def to_graph(builder, parent_id: nil, value: nil) + builder.add_node(node_id, name, Node::GATE) + builder.add_edge(parent_id, node_id, value) + + sink_nodes = branches.map do |condition_result, branch| + branch.to_graph(builder, parent_id: node_id, value: condition_result) + end + + builder.add_edge(node_id, "out", "default") + + sink_nodes.flatten + end + end + end + end +end diff --git a/lib/mars/rendering/graph/node.rb b/lib/mars/rendering/graph/node.rb new file mode 100644 index 0000000..1b11cb2 --- /dev/null +++ b/lib/mars/rendering/graph/node.rb @@ -0,0 +1,22 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + class Node + STEP = :step + OUTPUT = :output + INPUT = :input + GATE = :gate + + attr_reader :id, :name, :type + + def initialize(id, name, type) + @id = id + @name = name + @type = type + end + end + end + end +end diff --git a/lib/mars/rendering/graph/parallel_workflow.rb b/lib/mars/rendering/graph/parallel_workflow.rb new file mode 100644 index 0000000..0029aef --- /dev/null +++ b/lib/mars/rendering/graph/parallel_workflow.rb @@ -0,0 +1,24 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + module ParallelWorkflow + include Base + + def to_graph(builder, parent_id: nil, value: nil) + builder.add_node(aggregator.node_id, aggregator.name, Node::STEP) + + steps.each do |step| + sink_nodes = step.to_graph(builder, parent_id: parent_id, value: value) + sink_nodes.each do |sink_node| + aggregator.to_graph(builder, parent_id: sink_node) + end + end + + [aggregator.node_id] + end + end + end + end +end diff --git a/lib/mars/rendering/graph/sequential_workflow.rb b/lib/mars/rendering/graph/sequential_workflow.rb new file mode 100644 index 0000000..5c5dddf --- /dev/null +++ b/lib/mars/rendering/graph/sequential_workflow.rb @@ -0,0 +1,22 @@ +# frozen_string_literal: true + +module Mars + module Rendering + module Graph + module SequentialWorkflow + include Base + + def to_graph(builder, parent_id: nil, value: nil) + sink_nodes = [] + steps.each do |step| + sink_nodes = step.to_graph(builder, parent_id: parent_id, value: value) + value = nil # We don't want to pass the value to subsequent steps + parent_id = step.node_id + end + + sink_nodes.flatten + end + end + end + end +end diff --git a/lib/mars/rendering/mermaid.rb b/lib/mars/rendering/mermaid.rb index 5ca38e7..211974b 100644 --- a/lib/mars/rendering/mermaid.rb +++ b/lib/mars/rendering/mermaid.rb @@ -2,27 +2,58 @@ module Mars module Rendering - module Mermaid - def self.include_extensions - Mars::Agent.include(Agent) - Mars::Exit.include(Exit) - Mars::Gate.include(Gate) - Mars::Workflows::Sequential.include(SequentialWorkflow) - Mars::Workflows::Parallel.include(ParallelWorkflow) - Mars::Aggregator.include(Aggregator) + class Mermaid + attr_reader :obj, :graph, :nodes + + def initialize(obj) + @obj = obj + @graph, @nodes = obj.build_graph end - def self.render(obj, options = {}) + def render(options = {}) direction = options.fetch(:direction, "LR") + mermaid = graph_mermaid.join("\n") <<~MERMAID ```mermaid flowchart #{direction} - In(("In")) --> - #{obj.to_mermaid} + #{mermaid} ``` MERMAID end + + def graph_mermaid + nodes_mermaid = nodes.keys.map { |node_id| "#{node_id}#{shape(node_id)}" } + edges_mermaid = [] + + graph.each do |from, tos| + tos.each do |to| + node_id, value = to + edges_mermaid << "#{from} -->#{edge_value(value)} #{node_id}" + end + end + + nodes_mermaid + edges_mermaid + end + + def shape(node_id) + node = nodes[node_id] + + case node.type + when Graph::Node::INPUT, Graph::Node::OUTPUT + "((#{node.name}))" + when Graph::Node::GATE + "{#{node.name}}" + else + "[#{node.name}]" + end + end + + def edge_value(value) + return "" unless value + + "|#{value}|" + end end end end diff --git a/lib/mars/rendering/mermaid/agent.rb b/lib/mars/rendering/mermaid/agent.rb deleted file mode 100644 index 3ccbe9b..0000000 --- a/lib/mars/rendering/mermaid/agent.rb +++ /dev/null @@ -1,19 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module Agent - include Base - - def to_mermaid - "#{sanitized_name}[\"#{name}\"]" - end - - def can_end_workflow? - true - end - end - end - end -end diff --git a/lib/mars/rendering/mermaid/aggregator.rb b/lib/mars/rendering/mermaid/aggregator.rb deleted file mode 100644 index bf110dd..0000000 --- a/lib/mars/rendering/mermaid/aggregator.rb +++ /dev/null @@ -1,20 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module Aggregator - include Base - - # TODO - def to_mermaid - raise NotImplementedError - end - - def can_end_workflow? - true - end - end - end - end -end diff --git a/lib/mars/rendering/mermaid/base.rb b/lib/mars/rendering/mermaid/base.rb deleted file mode 100644 index 074d842..0000000 --- a/lib/mars/rendering/mermaid/base.rb +++ /dev/null @@ -1,17 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module Base - def sanitized_name - name.to_s.gsub(/[^a-zA-Z0-9_]/, "_") - end - - def can_end_workflow? - false - end - end - end - end -end diff --git a/lib/mars/rendering/mermaid/exit.rb b/lib/mars/rendering/mermaid/exit.rb deleted file mode 100644 index b7d3186..0000000 --- a/lib/mars/rendering/mermaid/exit.rb +++ /dev/null @@ -1,15 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module Exit - include Base - - def to_mermaid - "exit((#{name}))" - end - end - end - end -end diff --git a/lib/mars/rendering/mermaid/gate.rb b/lib/mars/rendering/mermaid/gate.rb deleted file mode 100644 index e475e43..0000000 --- a/lib/mars/rendering/mermaid/gate.rb +++ /dev/null @@ -1,28 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module Gate - include Base - - def to_mermaid - gate_id = sanitized_name - mermaid = ["#{gate_id}{\"#{name}\"}"] - - # Add edges for each branch - branches.each do |condition_result, branch| - branch_mermaid = branch.to_mermaid - mermaid << "#{gate_id} -->|#{condition_result}| #{branch_mermaid}" - end - - # Add the default exit path - default_mermaid = branches.default.to_mermaid - mermaid << "#{gate_id} -->|default| #{default_mermaid}" - - mermaid.join("\n") - end - end - end - end -end diff --git a/lib/mars/rendering/mermaid/parallel_workflow.rb b/lib/mars/rendering/mermaid/parallel_workflow.rb deleted file mode 100644 index 363564e..0000000 --- a/lib/mars/rendering/mermaid/parallel_workflow.rb +++ /dev/null @@ -1,16 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module ParallelWorkflow - include Base - - # TODO - def to_mermaid - raise NotImplementedError - end - end - end - end -end diff --git a/lib/mars/rendering/mermaid/sequential_workflow.rb b/lib/mars/rendering/mermaid/sequential_workflow.rb deleted file mode 100644 index e913433..0000000 --- a/lib/mars/rendering/mermaid/sequential_workflow.rb +++ /dev/null @@ -1,41 +0,0 @@ -# frozen_string_literal: true - -module Mars - module Rendering - module Mermaid - module SequentialWorkflow - include Base - - def to_mermaid - return "" if steps.empty? - - mermaid = steps.map.with_index do |current_step, index| - step_to_mermaid(current_step, index) - end - - mermaid.flatten.join("\n") - end - - private - - def step_to_mermaid(current_step, index) - current_id = current_step.sanitized_name - next_node = next_node(current_step, index) - - step_mermaid = [current_step.to_mermaid] - step_mermaid << "#{current_id} --> #{next_node}" if next_node - - step_mermaid - end - - def next_node(step, index) - if index < steps.length - 1 - steps[index + 1].sanitized_name - elsif step.can_end_workflow? - "Out((\"Out\"))" - end - end - end - end - end -end diff --git a/lib/mars/workflows/parallel.rb b/lib/mars/workflows/parallel.rb index 20ac9cc..6332827 100644 --- a/lib/mars/workflows/parallel.rb +++ b/lib/mars/workflows/parallel.rb @@ -8,7 +8,7 @@ class Parallel < Runnable def initialize(name, steps:, aggregator: nil) @name = name @steps = steps - @aggregator = aggregator || Aggregator.new + @aggregator = aggregator || Aggregator.new("#{name} Aggregator") end def run(input)