Add examples, mermaid diagram and exit step

examples/simple_workflow/diagram.md

@@ -0,0 +1,12 @@
+```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))
+```

examples/simple_workflow/generator.rb

@@ -0,0 +1,37 @@
+#!/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")
+
+# Create the success workflow (LLM 2 -> LLM 3)
+success_workflow = Mars::Workflows::Sequential.new(
+  "Success workflow",
+  steps: [llm2, llm3]
+)
+
+# Create the gate that decides between exit or continue
+gate = Mars::Gate.new(
+  name: "Gate",
+  condition: ->(input) { input[:result] },
+  branches: {
+    success: success_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.render(main_workflow)
+File.write("examples/simple_workflow/diagram.md", diagram)
+puts "Simple workflow diagram saved to: examples/simple_workflow/diagram.md"

lib/mars.rb

@@ -8,3 +8,5 @@
 module Mars
   class Error < StandardError; end
 end
+
+Mars::Rendering::Mermaid.include_extensions

lib/mars/agent.rb

@@ -2,7 +2,9 @@
 
 module Mars
   class Agent < Runnable
-    def initialize(name:, options:, tools: [], schema: nil)
+    attr_reader :name
+
+    def initialize(name:, options: {}, tools: [], schema: nil)
       @name = name
       @tools = Array(tools)
       @schema = schema
@@ -15,7 +17,7 @@ def run(input)
 
     private
 
-    attr_reader :name, :tools, :schema, :options
+    attr_reader :tools, :schema, :options
 
     def chat
       @chat ||= RubyLLM::Chat.new(**options).with_tools(tools).with_schema(schema)

lib/mars/exit.rb

@@ -0,0 +1,15 @@
+# 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

lib/mars/gate.rb

@@ -2,22 +2,22 @@
 
 module Mars
   class Gate < Runnable
+    attr_reader :name
+
     def initialize(name:, condition:, branches:)
       @name = name
       @condition = condition
-      @branches = branches
+      @branches = Hash.new(Exit.new).merge(branches)
     end
 
     def run(input)
       result = condition.call(input)
 
-      raise "Invalid condition result: #{result}" unless branches.key?(result)
-
       branches[result].run(input)
     end
 
     private
 
-    attr_reader :name, :condition, :branches
+    attr_reader :condition, :branches
   end
 end

lib/mars/rendering/mermaid.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+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)
+      end
+
+      def self.render(obj, options = {})
+        direction = options.fetch(:direction, "LR")
+
+        <<~MERMAID
+          ```mermaid
+          flowchart #{direction}
+          In(("In")) -->
+          #{obj.to_mermaid}
+          ```
+        MERMAID
+      end
+    end
+  end
+end

lib/mars/rendering/mermaid/agent.rb

@@ -0,0 +1,19 @@
+# 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

lib/mars/rendering/mermaid/base.rb

@@ -0,0 +1,17 @@
+# 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

lib/mars/rendering/mermaid/exit.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Mars
+  module Rendering
+    module Mermaid
+      module Exit
+        include Base
+
+        def to_mermaid
+          "exit((#{name}))"
+        end
+      end
+    end
+  end
+end

lib/mars/rendering/mermaid/gate.rb

@@ -0,0 +1,28 @@
+# 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

lib/mars/rendering/mermaid/sequential_workflow.rb

@@ -0,0 +1,41 @@
+# 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

lib/mars/workflows/sequential.rb

@@ -3,6 +3,8 @@
 module Mars
   module Workflows
     class Sequential < Runnable
+      attr_reader :name
+
       def initialize(name, steps:)
         @name = name
         @steps = steps
@@ -15,6 +17,10 @@ def run(input)
 
         input
       end
+
+      private
+
+      attr_reader :steps
     end
   end
 end