mars_rb 0.1.0 → 1.0.0

data/examples/simple_workflow/generator.rb

@@ -0,0 +1,48 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require_relative "../../lib/mars"
+
+# Define the LLMs
+class Agent1 < MARS::AgentStep
+end
+
+class Agent2 < MARS::AgentStep
+end
+
+class Agent3 < MARS::AgentStep
+end
+
+class Agent4 < MARS::AgentStep
+end
+
+# Create the LLMs
+llm1 = Agent1.new
+llm2 = Agent2.new
+llm3 = Agent3.new
+llm4 = Agent4.new
+
+# Create the failure workflow (LLM 3)
+failure_workflow = MARS::Workflows::Sequential.new(
+  "Failure workflow",
+  steps: [llm4]
+)
+
+# Create the gate that decides between exit or continue
+gate = MARS::Gate.new(
+  check: ->(input) { input[:result] },
+  fallbacks: {
+    failure: failure_workflow
+  }
+)
+
+# Create the main workflow: LLM 1 -> Gate
+main_workflow = MARS::Workflows::Sequential.new(
+  "Main Pipeline",
+  steps: [llm1, gate, llm2, llm3]
+)
+
+# Generate and save the diagram
+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"

data/lib/mars/agent_step.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module MARS
+  class AgentStep < Runnable
+    class << self
+      def agent(klass = nil)
+        klass ? @agent_class = klass : @agent_class
+      end
+    end
+
+    def run(context)
+      self.class.agent.new.ask(context.current_input).content
+    end
+  end
+end

data/lib/mars/aggregator.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module MARS
+  class Aggregator < Runnable
+    attr_reader :operation
+
+    def initialize(name = "Aggregator", operation: nil, **kwargs)
+      super(name: name, **kwargs)
+
+      @operation = operation || ->(inputs) { inputs }
+    end
+
+    def run(context)
+      context = ensure_context(context)
+      operation.call(context.current_input)
+    end
+
+    private
+
+    def ensure_context(input)
+      input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
+    end
+  end
+end

data/lib/mars/execution_context.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MARS
+  class ExecutionContext
+    attr_reader :outputs, :global_state
+    attr_accessor :current_input
+
+    def initialize(input: nil, global_state: {})
+      @current_input = input
+      @outputs = {}
+      @global_state = global_state
+    end
+
+    def [](step_name)
+      outputs[step_name.to_sym]
+    end
+
+    def record(step_name, output)
+      @outputs[step_name.to_sym] = output
+      @current_input = output
+    end
+
+    def fork(input: current_input, state: {})
+      self.class.new(input: input, global_state: global_state.merge(state))
+    end
+
+    def merge(child_contexts)
+      child_contexts.each do |child|
+        @outputs.merge!(child.outputs)
+      end
+
+      self
+    end
+  end
+end

data/lib/mars/formatter.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module MARS
+  class Formatter
+    def format_input(context)
+      context.current_input
+    end
+
+    def format_output(output)
+      output
+    end
+  end
+end

data/lib/mars/gate.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module MARS
+  class Gate < Runnable
+    class << self
+      def check(&block)
+        @check_block = block
+      end
+
+      attr_reader :check_block
+
+      def fallback(key, runnable)
+        fallbacks_map[key] = runnable
+      end
+
+      def fallbacks_map
+        @fallbacks_map ||= {}
+      end
+    end
+
+    def initialize(name = "Gate", check: nil, fallbacks: nil, **kwargs)
+      super(name: name, **kwargs)
+
+      @check = check || self.class.check_block
+      @fallbacks = fallbacks || self.class.fallbacks_map
+    end
+
+    def run(context)
+      context = ensure_context(context)
+      result = check.call(context)
+
+      return context if result.nil? || result == :default
+
+      branch = fallbacks[result]
+      raise ArgumentError, "No fallback registered for #{result.inspect}" unless branch
+
+      resolve_branch(branch).run(context)
+    end
+
+    private
+
+    attr_reader :check, :fallbacks
+
+    def resolve_branch(branch)
+      branch.is_a?(Class) ? branch.new : branch
+    end
+
+    def ensure_context(input)
+      input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
+    end
+  end
+end

data/lib/mars/hooks.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MARS
+  module Hooks
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      def before_run(&block)
+        before_run_hooks << block
+      end
+
+      def after_run(&block)
+        after_run_hooks << block
+      end
+
+      def before_run_hooks
+        @before_run_hooks ||= []
+      end
+
+      def after_run_hooks
+        @after_run_hooks ||= []
+      end
+    end
+
+    def run_before_hooks(context)
+      self.class.before_run_hooks.each { |hook| hook.call(context, self) }
+    end
+
+    def run_after_hooks(context, result)
+      self.class.after_run_hooks.each { |hook| hook.call(context, result, self) }
+    end
+  end
+end

data/lib/mars/rendering/graph/agent_step.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    module Graph
+      module AgentStep
+        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

data/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

data/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, builder.subgraphs]
+        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

data/lib/mars/rendering/graph/builder.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    module Graph
+      class Builder
+        attr_reader :adjacency, :nodes, :subgraphs
+
+        def initialize
+          @adjacency = Hash.new { |h, k| h[k] = [] }
+          @nodes = {}
+          @subgraphs = {}
+        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
+
+        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
+end

data/lib/mars/rendering/graph/gate.rb

@@ -0,0 +1,22 @@
+# 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 = fallbacks.map do |fallback_key, branch|
+            branch.to_graph(builder, parent_id: node_id, value: fallback_key)
+          end
+
+          sink_nodes.flatten
+        end
+      end
+    end
+  end
+end

data/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

data/lib/mars/rendering/graph/parallel_workflow.rb

@@ -0,0 +1,34 @@
+# 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_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
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    module Graph
+      module Runnable
+        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

data/lib/mars/rendering/graph/sequential_workflow.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    module Graph
+      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
+
+            builder.add_node_to_subgraph(node_id, step.node_id)
+
+            sink_nodes.each { |sink_node| builder.add_node_to_subgraph(node_id, sink_node) }
+          end
+
+          [parent_id, value, sink_nodes]
+        end
+      end
+    end
+  end
+end

data/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

data/lib/mars/rendering/graph.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    module Graph
+      def self.include_extensions
+        MARS::Runnable.include(Runnable)
+        MARS::AgentStep.include(AgentStep)
+        MARS::Gate.include(Gate)
+        MARS::Workflows::Sequential.include(SequentialWorkflow)
+        MARS::Workflows::Parallel.include(ParallelWorkflow)
+        MARS::Aggregator.include(Aggregator)
+      end
+    end
+  end
+end

data/lib/mars/rendering/html.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    class Html
+      MERMAID_CDN = "https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"
+
+      def initialize(obj)
+        @mermaid = Mermaid.new(obj)
+      end
+
+      def render
+        diagram = @mermaid.graph_mermaid.join("\n")
+        direction = "LR"
+
+        <<~HTML
+          <!DOCTYPE html>
+          <html lang="en">
+          <head>
+            <meta charset="UTF-8">
+            <meta name="viewport" content="width=device-width, initial-scale=1.0">
+            <title>MARS Workflow</title>
+            <script src="#{MERMAID_CDN}"></script>
+          </head>
+          <body>
+            <pre class="mermaid">
+          flowchart #{direction}
+          #{diagram}
+            </pre>
+            <script>mermaid.initialize({ startOnLoad: true });</script>
+          </body>
+          </html>
+        HTML
+      end
+    end
+  end
+end

data/lib/mars/rendering/mermaid.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    class Mermaid
+      attr_reader :obj, :graph, :nodes, :subgraphs
+
+      def initialize(obj)
+        @obj = obj
+        @graph, @nodes, @subgraphs = obj.build_graph
+      end
+
+      def render(options = {})
+        direction = options.fetch(:direction, "LR")
+        mermaid = graph_mermaid.join("\n")
+
+        <<~MERMAID
+          ```mermaid
+          flowchart #{direction}
+          #{mermaid}
+          ```
+        MERMAID
+      end
+
+      def graph_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 << "#{from} -->#{edge_value(value)} #{node_id}"
+          end
+        end
+        edges
+      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

data/lib/mars/runnable.rb

@@ -1,8 +1,32 @@
 # frozen_string_literal: true
 
-module Mars
+module MARS
   class Runnable
-    def run(input)
+    include Hooks
+
+    attr_reader :name, :formatter
+    attr_accessor :state
+
+    class << self
+      def step_name
+        return @step_name if defined?(@step_name)
+        return unless name
+
+        name.split("::").last.gsub(/([a-z])([A-Z])/, '\1_\2').downcase
+      end
+
+      def formatter(klass = nil)
+        klass ? @formatter_class = klass : @formatter_class
+      end
+    end
+
+    def initialize(name: self.class.step_name, state: {}, formatter: nil)
+      @name = name
+      @state = state
+      @formatter = formatter || self.class.formatter&.new || Formatter.new
+    end
+
+    def run(context)
       raise NotImplementedError
     end
   end

data/lib/mars/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
-module Mars
-  VERSION = "0.1.0"
+module MARS
+  VERSION = "1.0.0"
 end

data/lib/mars/workflows/aggregate_error.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module MARS
+  module Workflows
+    class AggregateError < StandardError
+      attr_reader :errors
+
+      def initialize(errors)
+        @errors = errors
+        super(errors.map { |error| "#{error[:step_name]}: #{error[:error].message}" }.join("\n"))
+      end
+    end
+  end
+end