Add subgraph rendering (#21)

* Add subgraph rendering

* Fixes

* More fixes

Author: PerezIgnacio

.rubocop.yml

@@ -19,6 +19,11 @@ Metrics/ParameterLists:
   Exclude:
   - "lib/mars/agent.rb"
 
+Metrics/AbcSize:
+  Enabled: true
+  Exclude:
+  - "lib/mars/rendering/graph/parallel_workflow.rb"
+
 Style/Documentation:
   Enabled: false
 

examples/complex_llm_workflow/diagram.md

@@ -8,6 +8,17 @@ parallel_workflow_aggregator[Parallel workflow Aggregator]
 agent2[Agent2]
 agent3[Agent3]
 agent4[Agent4]
+subgraph parallel_workflow["Parallel workflow"]
+  agent2
+  agent3
+  agent4
+end
+subgraph sequential_workflow["Sequential workflow"]
+  agent1
+  gate
+  parallel_workflow
+  parallel_workflow_aggregator
+end
 in --> agent1
 agent1 --> gate
 gate -->|failure| out

examples/complex_workflow/diagram.md

@@ -10,6 +10,26 @@ agent2[Agent2]
 agent3[Agent3]
 parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator]
 agent5[Agent5]
+subgraph parallel_workflow_2["Parallel workflow 2"]
+  sequential_workflow
+  agent5
+end
+subgraph parallel_workflow["Parallel workflow"]
+  agent2
+  agent3
+end
+subgraph sequential_workflow["Sequential workflow"]
+  agent4
+  parallel_workflow
+  parallel_workflow_aggregator
+end
+subgraph main_pipeline["Main Pipeline"]
+  agent1
+  gate
+  parallel_workflow_aggregator
+  parallel_workflow_2
+  parallel_workflow_2_aggregator
+end
 in --> agent1
 agent1 --> gate
 gate -->|warning| agent4

examples/parallel_workflow/diagram.md

@@ -6,6 +6,11 @@ aggregator[Aggregator]
 agent1[Agent1]
 agent2[Agent2]
 agent3[Agent3]
+subgraph parallel_workflow["Parallel workflow"]
+  agent1
+  agent2
+  agent3
+end
 in --> agent1
 in --> agent2
 in --> agent3

lib/mars/rendering/graph/base.rb

@@ -14,7 +14,7 @@ def build_graph(builder = Mars::Rendering::Graph::Builder.new)
             builder.add_edge(sink_node, "out")
           end
 
-          [builder.adjacency, builder.nodes]
+          [builder.adjacency, builder.nodes, builder.subgraphs]
         end
 
         def node_id

lib/mars/rendering/graph/builder.rb

@@ -4,11 +4,12 @@ module Mars
   module Rendering
     module Graph
       class Builder
-        attr_reader :adjacency, :nodes
+        attr_reader :adjacency, :nodes, :subgraphs
 
         def initialize
           @adjacency = Hash.new { |h, k| h[k] = [] }
           @nodes = {}
+          @subgraphs = {}
         end
 
         def add_edge(from, to, value = nil)
@@ -24,6 +25,18 @@ def add_node(id, value, type)
 
           nodes[id] = Node.new(id, value, type)
         end
+
+        def add_subgraph(id, name)
+          return if subgraphs.key?(id)
+
+          subgraphs[id] = Subgraph.new(id, name, [])
+        end
+
+        def add_node_to_subgraph(id, node_id)
+          return if subgraphs[id]&.nodes&.include?(node_id)
+
+          subgraphs[id].nodes << node_id
+        end
       end
     end
   end

lib/mars/rendering/graph/parallel_workflow.rb

@@ -7,16 +7,26 @@ module ParallelWorkflow
         include Base
 
         def to_graph(builder, parent_id: nil, value: nil)
+          builder.add_subgraph(node_id, name) if steps.any?
           builder.add_node(aggregator.node_id, aggregator.name, Node::STEP)
 
+          build_steps_graph(builder, parent_id, value)
+
+          [aggregator.node_id]
+        end
+
+        private
+
+        def build_steps_graph(builder, parent_id, value)
           steps.each do |step|
             sink_nodes = step.to_graph(builder, parent_id: parent_id, value: value)
+
+            builder.add_node_to_subgraph(node_id, step.node_id)
+
             sink_nodes.each do |sink_node|
               aggregator.to_graph(builder, parent_id: sink_node)
             end
           end
-
-          [aggregator.node_id]
         end
       end
     end

lib/mars/rendering/graph/sequential_workflow.rb

@@ -7,16 +7,31 @@ module SequentialWorkflow
         include Base
 
         def to_graph(builder, parent_id: nil, value: nil)
+          builder.add_subgraph(node_id, name) if steps.any?
+
+          parent_id, value, sink_nodes = build_steps_graph(builder, parent_id, value)
+
+          builder.add_edge(parent_id, "out", value) if sink_nodes.empty?
+
+          sink_nodes.flatten
+        end
+
+        private
+
+        def build_steps_graph(builder, parent_id, value)
           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
 
-          builder.add_edge(parent_id, "out", value) if sink_nodes.empty?
+            builder.add_node_to_subgraph(node_id, step.node_id)
 
-          sink_nodes.flatten
+            sink_nodes.each { |sink_node| builder.add_node_to_subgraph(node_id, sink_node) }
+          end
+
+          [parent_id, value, sink_nodes]
         end
       end
     end

lib/mars/rendering/graph/subgraph.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Mars
+  module Rendering
+    module Graph
+      class Subgraph
+        attr_reader :id, :name, :nodes
+
+        def initialize(id, name, nodes)
+          @id = id
+          @name = name
+          @nodes = nodes
+        end
+      end
+    end
+  end
+end

lib/mars/rendering/mermaid.rb

@@ -3,11 +3,11 @@
 module Mars
   module Rendering
     class Mermaid
-      attr_reader :obj, :graph, :nodes
+      attr_reader :obj, :graph, :nodes, :subgraphs
 
       def initialize(obj)
         @obj = obj
-        @graph, @nodes = obj.build_graph
+        @graph, @nodes, @subgraphs = obj.build_graph
       end
 
       def render(options = {})
@@ -23,17 +23,29 @@ def render(options = {})
       end
 
       def graph_mermaid
-        nodes_mermaid = nodes.keys.map { |node_id| "#{node_id}#{shape(node_id)}" }
-        edges_mermaid = []
+        nodes_mermaid + subgraphs_mermaid + edges_mermaid
+      end
+
+      def nodes_mermaid
+        nodes.keys.map { |node_id| "#{node_id}#{shape(node_id)}" }
+      end
 
+      def subgraphs_mermaid
+        subgraphs.values.reverse.map do |subgraph|
+          node_names = subgraph.nodes
+          "subgraph #{subgraph.id}[\"#{subgraph.name}\"]\n  #{node_names.join("\n  ")}\nend"
+        end
+      end
+
+      def edges_mermaid
+        edges = []
         graph.each do |from, tos|
           tos.each do |to|
             node_id, value = to
-            edges_mermaid << "#{from} -->#{edge_value(value)} #{node_id}"
+            edges << "#{from} -->#{edge_value(value)} #{node_id}"
           end
         end
-
-        nodes_mermaid + edges_mermaid
+        edges
       end
 
       def shape(node_id)