Change rendering approach

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
+```

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"

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
+```

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"

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
 ```

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"

lib/mars.rb

@@ -9,4 +9,4 @@ module Mars
   class Error < StandardError; end
 end
 
-Mars::Rendering::Mermaid.include_extensions
+Mars::Rendering::Graph.include_extensions

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)

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

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

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

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

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

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

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

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

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