diff --git a/README.md b/README.md
index 0c14456..d0aecca 100644
--- a/README.md
+++ b/README.md
@@ -47,47 +47,53 @@ Here's a simple example to get you started:
```ruby
require 'mars'
-# Define agents
-class Agent1 < Mars::Agent
+# Define a RubyLLM agent
+class MyAgent < RubyLLM::Agent
+ model "gpt-4o"
+ instructions "You are a helpful assistant."
end
-class Agent2 < Mars::Agent
+# Wrap it in a MARS step
+class MyStep < MARS::AgentStep
+ agent MyAgent
end
-class Agent3 < Mars::Agent
-end
-
-# Create agents
-agent1 = Agent1.new
-agent2 = Agent2.new
-agent3 = Agent3.new
+# Create steps
+step1 = MyStep.new(name: "step1")
+step2 = MyStep.new(name: "step2")
+step3 = MyStep.new(name: "step3")
# Create a sequential workflow
-workflow = Mars::Workflows::Sequential.new(
+workflow = MARS::Workflows::Sequential.new(
"My First Workflow",
- steps: [agent1, agent2, agent3]
+ steps: [step1, step2, step3]
)
# Run the workflow
-result = workflow.run("Your input here")
+context = workflow.run("Your input here")
+context.current_input # final output
+context[:step1] # access any step's output by name
```
## Core Concepts
-### Agents
+### Agent Steps
-Agents are the basic building blocks of MARS. They represent individual units of work:
+Agent steps are the basic building blocks of MARS. They wrap a `RubyLLM::Agent` subclass for workflow orchestration:
```ruby
-class CustomAgent < Mars::Agent
- def system_prompt
- "You are a helpful assistant"
- end
+class ResearcherAgent < RubyLLM::Agent
+ model "gpt-4o"
+ instructions "You research topics thoroughly."
+ tools WebSearch
+ schema OutputSchema
end
-agent = CustomAgent.new(
- options: { model: "gpt-4o" }
-)
+class ResearcherStep < MARS::AgentStep
+ agent ResearcherAgent
+end
+
+step = ResearcherStep.new(name: "researcher")
```
### Sequential Workflows
@@ -95,7 +101,7 @@ agent = CustomAgent.new(
Execute agents one after another, passing outputs as inputs:
```ruby
-sequential = Mars::Workflows::Sequential.new(
+sequential = MARS::Workflows::Sequential.new(
"Sequential Pipeline",
steps: [agent1, agent2, agent3]
)
@@ -106,12 +112,12 @@ sequential = Mars::Workflows::Sequential.new(
Run multiple agents concurrently and aggregate their results:
```ruby
-aggregator = Mars::Aggregator.new(
+aggregator = MARS::Aggregator.new(
"Results Aggregator",
operation: lambda { |results| results.join(", ") }
)
-parallel = Mars::Workflows::Parallel.new(
+parallel = MARS::Workflows::Parallel.new(
"Parallel Pipeline",
steps: [agent1, agent2, agent3],
aggregator: aggregator
@@ -123,7 +129,7 @@ parallel = Mars::Workflows::Parallel.new(
Create conditional branching in your workflows:
```ruby
-gate = Mars::Gate.new(
+gate = MARS::Gate.new(
"Decision Gate",
condition: ->(input) { input[:score] > 0.5 ? :success : :failure },
branches: {
@@ -138,7 +144,7 @@ gate = Mars::Gate.new(
Generate Mermaid diagrams to visualize your workflows:
```ruby
-diagram = Mars::Rendering::Mermaid.new(workflow).render
+diagram = MARS::Rendering::Mermaid.new(workflow).render
File.write("workflow_diagram.md", diagram)
```
@@ -197,7 +203,7 @@ The gem is available as open source under the terms of the [MIT License](https:/
## Code of Conduct
-Everyone interacting in the Mars project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/rootstrap/mars/blob/main/CODE_OF_CONDUCT.md).
+Everyone interacting in the MARS project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/rootstrap/mars/blob/main/CODE_OF_CONDUCT.md).
## Credits
diff --git a/examples/complex_llm_workflow/generator.rb b/examples/complex_llm_workflow/generator.rb
index 4ea2c77..0bfe4b6 100755
--- a/examples/complex_llm_workflow/generator.rb
+++ b/examples/complex_llm_workflow/generator.rb
@@ -37,48 +37,52 @@ def execute(latitude:, longitude:)
end
end
-# Define LLMs
-class Agent1 < MARS::Agent
- def system_prompt
- "You are a helpful assistant that can answer questions.
- When asked about a country, only answer with its name."
- end
+# Define RubyLLM agents
+class CountryAgent < RubyLLM::Agent
+ model "gpt-4o"
+ instructions "You are a helpful assistant that can answer questions. " \
+ "When asked about a country, only answer with its name."
end
-class Agent2 < MARS::Agent
- def system_prompt
- "You are a helpful assistant that can answer questions and help with tasks.
- Return information about the typical food of the country."
- end
+class FoodAgent < RubyLLM::Agent
+ model "gpt-4o"
+ instructions "You are a helpful assistant. Return information about the typical food of the country."
end
-class Agent3 < MARS::Agent
- def system_prompt
- "You are a helpful assistant that can answer questions and help with tasks.
- Return information about the popular sports of the country."
- end
+class SportsAgent < RubyLLM::Agent
+ model "gpt-4o"
+ instructions "You are a helpful assistant. Return information about the popular sports of the country."
+ schema SportsSchema
+end
- def schema
- SportsSchema.new
- end
+class WeatherAgent < RubyLLM::Agent
+ model "gpt-4o"
+ instructions "You are a helpful assistant. Return the current weather of the country's capital."
+ tools Weather
end
-class Agent4 < MARS::Agent
- def system_prompt
- "You are a helpful assistant that can answer questions and help with tasks.
- Return the current weather of the country's capital."
- end
+# Define MARS steps wrapping RubyLLM agents
+class CountryStep < MARS::AgentStep
+ agent CountryAgent
+end
- def tools
- [Weather.new]
- end
+class FoodStep < MARS::AgentStep
+ agent FoodAgent
+end
+
+class SportsStep < MARS::AgentStep
+ agent SportsAgent
+end
+
+class WeatherStep < MARS::AgentStep
+ agent WeatherAgent
end
-# Create the LLMs
-llm1 = Agent1.new(options: { model: "gpt-4o" })
-llm2 = Agent2.new(options: { model: "gpt-4o" })
-llm3 = Agent3.new(options: { model: "gpt-4o" })
-llm4 = Agent4.new(options: { model: "gpt-4o" })
+# Create the steps
+llm1 = CountryStep.new(name: "Country")
+llm2 = FoodStep.new(name: "Food")
+llm3 = SportsStep.new(name: "Sports")
+llm4 = WeatherStep.new(name: "Weather")
parallel_workflow = MARS::Workflows::Parallel.new(
"Parallel workflow",
diff --git a/examples/complex_workflow/diagram.html b/examples/complex_workflow/diagram.html
new file mode 100644
index 0000000..53b7070
--- /dev/null
+++ b/examples/complex_workflow/diagram.html
@@ -0,0 +1,64 @@
+
+
+
+
+
+ Main Pipeline
+
+
+
+ Main Pipeline
+
+
+
+
diff --git a/examples/complex_workflow/diagram.md b/examples/complex_workflow/diagram.md
index a424141..dd8fc85 100644
--- a/examples/complex_workflow/diagram.md
+++ b/examples/complex_workflow/diagram.md
@@ -2,32 +2,29 @@
flowchart LR
in((In))
out((Out))
-agent1[Agent1]
+agent1[agent1]
gate{Gate}
-agent4[Agent4]
+agent4[agent4]
parallel_workflow_aggregator[Parallel workflow Aggregator]
-agent2[Agent2]
-agent3[Agent3]
+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
+agent5[agent5]
subgraph main_pipeline["Main Pipeline"]
agent1
gate
parallel_workflow_aggregator
- parallel_workflow_2
+ subgraph parallel_workflow_2["Parallel workflow 2"]
+ subgraph sequential_workflow["Sequential workflow"]
+ agent4
+ subgraph parallel_workflow["Parallel workflow"]
+ agent2
+ agent3
+ end
+ parallel_workflow_aggregator
+ end
+ agent5
+ end
parallel_workflow_2_aggregator
end
in --> agent1
diff --git a/examples/complex_workflow/generator.rb b/examples/complex_workflow/generator.rb
index 62adc6b..d32a2fa 100755
--- a/examples/complex_workflow/generator.rb
+++ b/examples/complex_workflow/generator.rb
@@ -4,19 +4,19 @@
require_relative "../../lib/mars"
# Define LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
end
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
end
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
end
-class Agent4 < MARS::Agent
+class Agent4 < MARS::AgentStep
end
-class Agent5 < MARS::Agent
+class Agent5 < MARS::AgentStep
end
# Create the LLMs
diff --git a/examples/parallel_workflow/generator.rb b/examples/parallel_workflow/generator.rb
index 795c2c5..66378bd 100755
--- a/examples/parallel_workflow/generator.rb
+++ b/examples/parallel_workflow/generator.rb
@@ -4,13 +4,13 @@
require_relative "../../lib/mars"
# Define the LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
end
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
end
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
end
# Create the LLMs
diff --git a/examples/simple_workflow/generator.rb b/examples/simple_workflow/generator.rb
index 4dd01c4..b1a0351 100755
--- a/examples/simple_workflow/generator.rb
+++ b/examples/simple_workflow/generator.rb
@@ -4,13 +4,13 @@
require_relative "../../lib/mars"
# Define the LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
end
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
end
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
end
# Create the LLMs
diff --git a/lib/mars/agent.rb b/lib/mars/agent.rb
deleted file mode 100644
index edc7ed7..0000000
--- a/lib/mars/agent.rb
+++ /dev/null
@@ -1,48 +0,0 @@
-# frozen_string_literal: true
-
-module MARS
- class Agent < Runnable
- def initialize(options: {}, **kwargs)
- super(**kwargs)
-
- @options = options
- end
-
- def run(input)
- processed_input = before_run(input)
- response = chat.ask(processed_input).content
- after_run(response)
- end
-
- private
-
- attr_reader :options
-
- def chat
- @chat ||= RubyLLM::Chat.new(**options)
- .with_instructions(system_prompt)
- .with_tools(*tools)
- .with_schema(schema)
- end
-
- def before_run(input)
- input
- end
-
- def after_run(response)
- response
- end
-
- def system_prompt
- nil
- end
-
- def tools
- []
- end
-
- def schema
- nil
- end
- end
-end
diff --git a/lib/mars/agent_step.rb b/lib/mars/agent_step.rb
new file mode 100644
index 0000000..a519997
--- /dev/null
+++ b/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(input)
+ self.class.agent.new.ask(input).content
+ end
+ end
+end
diff --git a/lib/mars/aggregator.rb b/lib/mars/aggregator.rb
index 0866e9f..78f5752 100644
--- a/lib/mars/aggregator.rb
+++ b/lib/mars/aggregator.rb
@@ -2,17 +2,20 @@
module MARS
class Aggregator < Runnable
- attr_reader :name, :operation
+ attr_reader :operation
def initialize(name = "Aggregator", operation: nil, **kwargs)
- super(**kwargs)
+ super(name: name, **kwargs)
- @name = name
@operation = operation || ->(inputs) { inputs }
end
def run(inputs)
- operation.call(inputs)
+ if inputs.is_a?(ExecutionContext)
+ operation.call(inputs.outputs)
+ else
+ operation.call(inputs)
+ end
end
end
end
diff --git a/lib/mars/execution_context.rb b/lib/mars/execution_context.rb
new file mode 100644
index 0000000..8d7e1b8
--- /dev/null
+++ b/lib/mars/execution_context.rb
@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module MARS
+ class ExecutionContext
+ attr_reader :current_input, :outputs, :global_state
+
+ def initialize(input: nil, global_state: {})
+ @current_input = input
+ @outputs = {}
+ @global_state = global_state
+ end
+
+ def [](step_name)
+ outputs[step_name]
+ end
+
+ def record(step_name, output)
+ @outputs[step_name.to_sym] = output
+ @current_input = output
+ end
+
+ def fork(input: current_input)
+ self.class.new(input: input, global_state: global_state)
+ end
+
+ def merge(child_contexts)
+ child_contexts.each do |child|
+ @outputs.merge!(child.outputs)
+ end
+
+ self
+ end
+ end
+end
diff --git a/lib/mars/formatter.rb b/lib/mars/formatter.rb
new file mode 100644
index 0000000..f82fb9c
--- /dev/null
+++ b/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
diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb
index 21e407d..e416ab8 100644
--- a/lib/mars/gate.rb
+++ b/lib/mars/gate.rb
@@ -2,24 +2,44 @@
module MARS
class Gate < Runnable
- attr_reader :name
+ class << self
+ def condition(&block)
+ @condition_block = block
+ end
- def initialize(name = "Gate", condition:, branches:, **kwargs)
- super(**kwargs)
+ attr_reader :condition_block
- @name = name
- @condition = condition
- @branches = branches
+ def branch(key, runnable)
+ branches_map[key] = runnable
+ end
+
+ def branches_map
+ @branches_map ||= {}
+ end
+ end
+
+ def initialize(name = "Gate", condition: nil, branches: nil, **kwargs)
+ super(name: name, **kwargs)
+
+ @condition = condition || self.class.condition_block
+ @branches = branches || self.class.branches_map
end
def run(input)
result = condition.call(input)
+ branch = branches[result]
- branches[result] || input
+ return input unless branch
+
+ resolve_branch(branch).run(input)
end
private
attr_reader :condition, :branches
+
+ def resolve_branch(branch)
+ branch.is_a?(Class) ? branch.new : branch
+ end
end
end
diff --git a/lib/mars/hooks.rb b/lib/mars/hooks.rb
new file mode 100644
index 0000000..df616ad
--- /dev/null
+++ b/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
diff --git a/lib/mars/rendering/graph.rb b/lib/mars/rendering/graph.rb
index 1a73d66..c453a81 100644
--- a/lib/mars/rendering/graph.rb
+++ b/lib/mars/rendering/graph.rb
@@ -4,7 +4,8 @@ module MARS
module Rendering
module Graph
def self.include_extensions
- MARS::Agent.include(Agent)
+ MARS::Runnable.include(Base)
+ MARS::AgentStep.include(AgentStep)
MARS::Gate.include(Gate)
MARS::Workflows::Sequential.include(SequentialWorkflow)
MARS::Workflows::Parallel.include(ParallelWorkflow)
diff --git a/lib/mars/rendering/graph/agent.rb b/lib/mars/rendering/graph/agent_step.rb
similarity index 81%
rename from lib/mars/rendering/graph/agent.rb
rename to lib/mars/rendering/graph/agent_step.rb
index 48f24d8..e170feb 100644
--- a/lib/mars/rendering/graph/agent.rb
+++ b/lib/mars/rendering/graph/agent_step.rb
@@ -3,7 +3,7 @@
module MARS
module Rendering
module Graph
- module Agent
+ module AgentStep
include Base
def to_graph(builder, parent_id: nil, value: nil)
@@ -12,10 +12,6 @@ def to_graph(builder, parent_id: nil, value: nil)
[node_id]
end
-
- def name
- self.class.name
- end
end
end
end
diff --git a/lib/mars/rendering/graph/base.rb b/lib/mars/rendering/graph/base.rb
index 2e70f54..d785470 100644
--- a/lib/mars/rendering/graph/base.rb
+++ b/lib/mars/rendering/graph/base.rb
@@ -21,6 +21,13 @@ def node_id
@node_id ||= sanitize(name)
end
+ 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
+
private
def sanitize(name)
diff --git a/lib/mars/rendering/html.rb b/lib/mars/rendering/html.rb
new file mode 100644
index 0000000..09c1297
--- /dev/null
+++ b/lib/mars/rendering/html.rb
@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module MARS
+ module Rendering
+ class Html
+ BEAUTIFUL_MERMAID_URL = "https://esm.sh/beautiful-mermaid@1"
+
+ attr_reader :obj
+
+ def initialize(obj)
+ @obj = obj
+ end
+
+ def render(options = {})
+ mermaid = Mermaid.new(obj)
+ diagram = mermaid.graph_mermaid.join("\n")
+ direction = options.fetch(:direction, "LR")
+ title = options.fetch(:title, obj.name)
+ theme = options.fetch(:theme, {})
+
+ build_html(title, direction, diagram, theme)
+ end
+
+ def write(path, options = {})
+ File.write(path, render(options))
+ end
+
+ private
+
+ def build_html(title, direction, diagram, theme)
+ <<~HTML
+
+
+
+
+
+ #{escape(title)}
+ #{head_style}
+
+
+ #{escape(title)}
+
+ #{render_script(direction, diagram, theme)}
+
+
+ HTML
+ end
+
+ def head_style
+ <<~STYLE.chomp
+
+ STYLE
+ end
+
+ def render_script(direction, diagram, theme)
+ theme_opts = theme_options(theme)
+
+ <<~SCRIPT.chomp
+
+ SCRIPT
+ end
+
+ def theme_options(theme)
+ return "" if theme.empty?
+
+ pairs = theme.map { |k, v| "#{k}: '#{escape_js(v.to_s)}'" }
+ ", { #{pairs.join(", ")} }"
+ end
+
+ def escape(text)
+ text.to_s.gsub("&", "&").gsub("<", "<").gsub(">", ">").gsub('"', """)
+ end
+
+ def escape_js(text)
+ text.to_s.gsub("\\", "\\\\\\\\").gsub("`", "\\`").gsub("$", "\\$")
+ end
+ end
+ end
+end
diff --git a/lib/mars/rendering/mermaid.rb b/lib/mars/rendering/mermaid.rb
index b9f2ac8..3f1914a 100644
--- a/lib/mars/rendering/mermaid.rb
+++ b/lib/mars/rendering/mermaid.rb
@@ -31,10 +31,8 @@ def nodes_mermaid
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
+ root_ids = subgraphs.keys - nested_subgraph_ids
+ root_ids.map { |id| render_subgraph(id) }
end
def edges_mermaid
@@ -66,6 +64,32 @@ def edge_value(value)
"|#{value}|"
end
+
+ private
+
+ def nested_subgraph_ids
+ ids = []
+ subgraphs.each_value do |sg|
+ sg.nodes.each { |n| ids << n if subgraphs.key?(n) }
+ end
+ ids
+ end
+
+ def render_subgraph(id, indent = "")
+ sg = subgraphs[id]
+ lines = ["#{indent}subgraph #{sg.id}[\"#{sg.name}\"]"]
+ sg.nodes.each { |node_id| lines << render_subgraph_node(node_id, indent) }
+ lines << "#{indent}end"
+ lines.join("\n")
+ end
+
+ def render_subgraph_node(node_id, indent)
+ if subgraphs.key?(node_id)
+ render_subgraph(node_id, "#{indent} ")
+ else
+ "#{indent} #{node_id}"
+ end
+ end
end
end
end
diff --git a/lib/mars/runnable.rb b/lib/mars/runnable.rb
index f7a4909..115baf1 100644
--- a/lib/mars/runnable.rb
+++ b/lib/mars/runnable.rb
@@ -2,10 +2,30 @@
module MARS
class Runnable
+ include Hooks
+
+ attr_reader :name, :formatter
attr_accessor :state
- def initialize(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
+
+ attr_writer :step_name
+
+ 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(input)
diff --git a/lib/mars/workflows/parallel.rb b/lib/mars/workflows/parallel.rb
index da65a27..b86912e 100644
--- a/lib/mars/workflows/parallel.rb
+++ b/lib/mars/workflows/parallel.rb
@@ -3,38 +3,50 @@
module MARS
module Workflows
class Parallel < Runnable
- attr_reader :name
-
def initialize(name, steps:, aggregator: nil, **kwargs)
- super(**kwargs)
+ super(name: name, **kwargs)
- @name = name
@steps = steps
@aggregator = aggregator || Aggregator.new("#{name} Aggregator")
end
def run(input)
- errors = []
- results = Async do |workflow|
- tasks = @steps.map do |step|
- workflow.async do
- step.run(input)
- rescue StandardError => e
- errors << { error: e, step_name: step.name }
- end
- end
+ context = input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
- tasks.map(&:wait)
- end.result
+ errors = []
+ child_contexts = run_steps_async(context, errors)
raise AggregateError, errors if errors.any?
- aggregator.run(results)
+ context.merge(child_contexts)
+ aggregator.run(context)
end
private
attr_reader :steps, :aggregator
+
+ def run_steps_async(context, errors)
+ Async do |workflow|
+ tasks = steps.map do |step|
+ workflow.async { run_step(context.fork, step, errors) }
+ end
+
+ tasks.map(&:wait)
+ end.result
+ end
+
+ def run_step(child, step, errors)
+ step.run_before_hooks(child)
+ step_input = step.formatter.format_input(child)
+ result = step.run(step_input)
+ formatted = step.formatter.format_output(result)
+ child.record(step.name, formatted)
+ step.run_after_hooks(child, formatted)
+ child
+ rescue StandardError => e
+ errors << { error: e, step_name: step.name }
+ end
end
end
end
diff --git a/lib/mars/workflows/sequential.rb b/lib/mars/workflows/sequential.rb
index 0db625e..e30e15d 100644
--- a/lib/mars/workflows/sequential.rb
+++ b/lib/mars/workflows/sequential.rb
@@ -3,33 +3,48 @@
module MARS
module Workflows
class Sequential < Runnable
- attr_reader :name
-
def initialize(name, steps:, **kwargs)
- super(**kwargs)
+ super(name: name, **kwargs)
- @name = name
@steps = steps
end
+ def self.build(name, **kwargs, &)
+ builder = Builder.new
+ builder.instance_eval(&)
+ new(name, steps: builder.steps, **kwargs)
+ end
+
def run(input)
- @steps.each do |step|
- result = step.run(input)
-
- if result.is_a?(Runnable)
- input = result.run(input)
- break
- else
- input = result
- end
+ context = input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
+
+ steps.each do |step|
+ step.run_before_hooks(context)
+ step_input = step.formatter.format_input(context)
+ result = step.run(step_input)
+ formatted = step.formatter.format_output(result)
+ context.record(step.name, formatted)
+ step.run_after_hooks(context, formatted)
end
- input
+ context
end
private
attr_reader :steps
+
+ class Builder
+ attr_reader :steps
+
+ def initialize
+ @steps = []
+ end
+
+ def step(runnable_class, **kwargs)
+ @steps << runnable_class.new(**kwargs)
+ end
+ end
end
end
end
diff --git a/mars.gemspec b/mars.gemspec
index 2e77793..55c3586 100644
--- a/mars.gemspec
+++ b/mars.gemspec
@@ -36,7 +36,7 @@ Gem::Specification.new do |spec|
# Uncomment to register a new dependency of your gem
spec.add_dependency "async", "~> 2.34"
- spec.add_dependency "ruby_llm", "~> 1.9"
+ spec.add_dependency "ruby_llm", "~> 1.12"
spec.add_dependency "zeitwerk", "~> 2.7"
# For more information and examples about making a new gem, check out our
diff --git a/spec/mars/agent_spec.rb b/spec/mars/agent_spec.rb
deleted file mode 100644
index 440f8da..0000000
--- a/spec/mars/agent_spec.rb
+++ /dev/null
@@ -1,64 +0,0 @@
-# frozen_string_literal: true
-
-RSpec.describe MARS::Agent do
- describe "#run" do
- subject(:run_agent) { agent.run("input text") }
-
- let(:agent) { described_class.new(options: { model: "test-model" }) }
- let(:mock_chat_instance) do
- instance_double("RubyLLM::Chat").tap do |mock|
- allow(mock).to receive_messages(with_tools: mock, with_schema: mock, with_instructions: mock,
- ask: mock_chat_response)
- end
- end
- let(:mock_chat_response) { instance_double("RubyLLM::Message", content: "response text") }
- let(:mock_chat_class) { class_double("RubyLLM::Chat", new: mock_chat_instance) }
-
- before do
- stub_const("RubyLLM::Chat", mock_chat_class)
- end
-
- it "initializes RubyLLM::Chat with provided options" do
- run_agent
-
- expect(mock_chat_class).to have_received(:new).with(model: "test-model")
- end
-
- context "when tools are provided" do
- let(:tools) { [proc { "tool1" }, proc { "tool2" }] }
- let(:agent_class) do
- Class.new(described_class) do
- def tools
- [proc { "tool1" }, proc { "tool2" }]
- end
- end
- end
-
- let(:agent) { agent_class.new }
-
- it "configures chat with tools" do
- run_agent
-
- expect(mock_chat_instance).to have_received(:with_tools).with(*tools)
- end
- end
-
- context "when schema is provided" do
- let(:agent_class) do
- Class.new(described_class) do
- def schema
- { type: "object" }
- end
- end
- end
-
- let(:agent) { agent_class.new }
-
- it "configures chat with schema" do
- run_agent
-
- expect(mock_chat_instance).to have_received(:with_schema).with({ type: "object" })
- end
- end
- end
-end
diff --git a/spec/mars/agent_step_spec.rb b/spec/mars/agent_step_spec.rb
new file mode 100644
index 0000000..18244a5
--- /dev/null
+++ b/spec/mars/agent_step_spec.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::AgentStep do
+ describe ".agent" do
+ it "stores and retrieves the agent class" do
+ agent_class = Class.new
+ step_class = Class.new(described_class) do
+ agent agent_class
+ end
+
+ expect(step_class.agent).to eq(agent_class)
+ end
+
+ it "returns nil when no agent class is set" do
+ step_class = Class.new(described_class)
+ expect(step_class.agent).to be_nil
+ end
+ end
+
+ describe "#run" do
+ let(:mock_agent_instance) do
+ instance_double("RubyLLM::Agent").tap do |mock|
+ allow(mock).to receive(:ask).and_return(instance_double("RubyLLM::Message", content: "agent response"))
+ end
+ end
+
+ let(:mock_agent_class) do
+ instance_double("Class").tap do |mock|
+ allow(mock).to receive(:new).and_return(mock_agent_instance)
+ end
+ end
+
+ let(:step_class) do
+ klass = mock_agent_class
+ Class.new(described_class) do
+ agent klass
+ end
+ end
+
+ it "creates a new agent instance and calls ask" do
+ step = step_class.new
+ result = step.run("hello")
+
+ expect(result).to eq("agent response")
+ expect(mock_agent_class).to have_received(:new)
+ expect(mock_agent_instance).to have_received(:ask).with("hello")
+ end
+ end
+
+ describe "inheritance" do
+ it "inherits from MARS::Runnable" do
+ expect(described_class.ancestors).to include(MARS::Runnable)
+ end
+
+ it "has access to name, formatter, and hooks from Runnable" do
+ step = described_class.new(name: "my_agent")
+ expect(step.name).to eq("my_agent")
+ expect(step.formatter).to be_a(MARS::Formatter)
+ end
+ end
+end
diff --git a/spec/mars/aggregator_spec.rb b/spec/mars/aggregator_spec.rb
index 294ec93..dd652da 100644
--- a/spec/mars/aggregator_spec.rb
+++ b/spec/mars/aggregator_spec.rb
@@ -2,7 +2,7 @@
RSpec.describe MARS::Aggregator do
describe "#run" do
- context "when called without a block" do
+ context "when called without an operation" do
let(:aggregator) { described_class.new }
it "returns the input as is" do
@@ -11,13 +11,28 @@
end
end
- context "when initialized with a block operation" do
+ context "when initialized with an operation" do
let(:aggregator) { described_class.new("Aggregator", operation: lambda(&:join)) }
- it "executes the block and returns its value" do
+ it "executes the operation and returns its value" do
result = aggregator.run(%w[a b c])
expect(result).to eq("abc")
end
end
+
+ context "when given an ExecutionContext" do
+ let(:aggregator) do
+ described_class.new("ContextAggregator", operation: ->(outputs) { outputs.values.join(", ") })
+ end
+
+ it "passes the context outputs to the operation" do
+ context = MARS::ExecutionContext.new(input: "query")
+ context.record(:step_a, "result_a")
+ context.record(:step_b, "result_b")
+
+ result = aggregator.run(context)
+ expect(result).to eq("result_a, result_b")
+ end
+ end
end
end
diff --git a/spec/mars/execution_context_spec.rb b/spec/mars/execution_context_spec.rb
new file mode 100644
index 0000000..d30d47d
--- /dev/null
+++ b/spec/mars/execution_context_spec.rb
@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::ExecutionContext do
+ describe "#current_input" do
+ it "returns the initial input" do
+ context = described_class.new(input: "query")
+ expect(context.current_input).to eq("query")
+ end
+
+ it "returns nil when no input is provided" do
+ context = described_class.new
+ expect(context.current_input).to be_nil
+ end
+ end
+
+ describe "#record" do
+ it "stores the output under the step name" do
+ context = described_class.new(input: "query")
+ context.record(:researcher, "research result")
+
+ expect(context[:researcher]).to eq("research result")
+ end
+
+ it "updates current_input to the recorded output" do
+ context = described_class.new(input: "query")
+ context.record(:researcher, "research result")
+
+ expect(context.current_input).to eq("research result")
+ end
+
+ it "tracks multiple step outputs" do
+ context = described_class.new(input: "query")
+ context.record(:researcher, "research result")
+ context.record(:summarizer, "summary")
+
+ expect(context.outputs).to eq({ researcher: "research result", summarizer: "summary" })
+ expect(context.current_input).to eq("summary")
+ end
+ end
+
+ describe "#[]" do
+ it "returns nil for unknown step names" do
+ context = described_class.new(input: "query")
+ expect(context[:unknown]).to be_nil
+ end
+ end
+
+ describe "#global_state" do
+ it "defaults to an empty hash" do
+ context = described_class.new(input: "query")
+ expect(context.global_state).to eq({})
+ end
+
+ it "accepts initial global state" do
+ context = described_class.new(input: "query", global_state: { user_id: 42 })
+ expect(context.global_state).to eq({ user_id: 42 })
+ end
+
+ it "is mutable" do
+ context = described_class.new(input: "query")
+ context.global_state[:key] = "value"
+
+ expect(context.global_state[:key]).to eq("value")
+ end
+ end
+
+ describe "#fork" do
+ it "creates a child context with the same current_input by default" do
+ context = described_class.new(input: "query")
+ child = context.fork
+
+ expect(child.current_input).to eq("query")
+ end
+
+ it "creates a child context with a custom input" do
+ context = described_class.new(input: "query")
+ child = context.fork(input: "custom")
+
+ expect(child.current_input).to eq("custom")
+ end
+
+ it "shares global_state with the parent" do
+ context = described_class.new(input: "query", global_state: { shared: true })
+ child = context.fork
+
+ child.global_state[:added_by_child] = true
+
+ expect(context.global_state[:added_by_child]).to be(true)
+ end
+
+ it "has independent outputs from the parent" do
+ context = described_class.new(input: "query")
+ context.record(:parent_step, "parent output")
+
+ child = context.fork
+ child.record(:child_step, "child output")
+
+ expect(context[:child_step]).to be_nil
+ expect(child[:parent_step]).to be_nil
+ end
+ end
+
+ describe "#merge" do
+ it "merges child outputs into the parent" do
+ context = described_class.new(input: "query")
+ context.record(:step1, "output1")
+
+ child1 = context.fork
+ child1.record(:branch_a, "result_a")
+
+ child2 = context.fork
+ child2.record(:branch_b, "result_b")
+
+ context.merge([child1, child2])
+
+ expect(context[:step1]).to eq("output1")
+ expect(context[:branch_a]).to eq("result_a")
+ expect(context[:branch_b]).to eq("result_b")
+ end
+ end
+end
diff --git a/spec/mars/formatter_spec.rb b/spec/mars/formatter_spec.rb
new file mode 100644
index 0000000..55b17f3
--- /dev/null
+++ b/spec/mars/formatter_spec.rb
@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::Formatter do
+ let(:formatter) { described_class.new }
+
+ describe "#format_input" do
+ it "returns the context's current_input" do
+ context = MARS::ExecutionContext.new(input: "hello")
+ expect(formatter.format_input(context)).to eq("hello")
+ end
+ end
+
+ describe "#format_output" do
+ it "returns the output unchanged" do
+ expect(formatter.format_output("result")).to eq("result")
+ end
+ end
+
+ describe "custom formatter" do
+ let(:custom_formatter_class) do
+ Class.new(described_class) do
+ def format_input(context)
+ context.current_input.upcase
+ end
+
+ def format_output(output)
+ "formatted: #{output}"
+ end
+ end
+ end
+
+ let(:custom_formatter) { custom_formatter_class.new }
+
+ it "can override format_input" do
+ context = MARS::ExecutionContext.new(input: "hello")
+ expect(custom_formatter.format_input(context)).to eq("HELLO")
+ end
+
+ it "can override format_output" do
+ expect(custom_formatter.format_output("result")).to eq("formatted: result")
+ end
+ end
+end
diff --git a/spec/mars/gate_spec.rb b/spec/mars/gate_spec.rb
index d8cf33d..0990e58 100644
--- a/spec/mars/gate_spec.rb
+++ b/spec/mars/gate_spec.rb
@@ -2,82 +2,120 @@
RSpec.describe MARS::Gate do
describe "#run" do
- let(:gate) { described_class.new("TestGate", condition: condition, branches: branches) }
+ context "with constructor-based configuration" do
+ let(:short_step) do
+ Class.new(MARS::Runnable) do
+ def run(input)
+ "short: #{input}"
+ end
+ end.new
+ end
- context "with simple boolean condition" do
- let(:condition) { ->(input) { input > 5 } }
- let(:false_branch) { instance_spy(MARS::Runnable) }
- let(:branches) { { false => false_branch } }
+ let(:long_step) do
+ Class.new(MARS::Runnable) do
+ def run(input)
+ "long: #{input}"
+ end
+ end.new
+ end
- it "returns the input when no branch matches" do
- result = gate.run(10)
- expect(result).to eq(10)
+ let(:gate) do
+ described_class.new(
+ "LengthGate",
+ condition: ->(input) { input.length > 5 ? :long : :short },
+ branches: { short: short_step, long: long_step }
+ )
end
- it "returns the false branch when condition is false" do
- result = gate.run(3)
+ it "executes the matched branch directly" do
+ expect(gate.run("hi")).to eq("short: hi")
+ end
- expect(result).to eq(false_branch)
+ it "executes the other branch for different input" do
+ expect(gate.run("longstring")).to eq("long: longstring")
end
- it "does not run the false branch when condition is false" do
- gate.run(3)
+ it "returns input when no branch matches" do
+ gate = described_class.new(
+ "NoMatch",
+ condition: ->(_input) { :unknown },
+ branches: { short: short_step }
+ )
- expect(false_branch).not_to have_received(:run)
+ expect(gate.run("hello")).to eq("hello")
end
end
- context "with string-based condition" do
- let(:condition) { ->(input) { input.length > 5 ? "long" : "short" } }
- let(:long_branch) { instance_spy(MARS::Runnable) }
- let(:short_branch) { instance_spy(MARS::Runnable) }
- let(:branches) { { "long" => long_branch, "short" => short_branch } }
-
- it "routes to long branch for long strings" do
- result = gate.run("longstring")
+ context "with class-level DSL" do
+ let(:short_step_class) do
+ Class.new(MARS::Runnable) do
+ def run(input)
+ "quick: #{input}"
+ end
+ end
+ end
- expect(result).to eq(long_branch)
+ let(:long_step_class) do
+ Class.new(MARS::Runnable) do
+ def run(input)
+ "deep: #{input}"
+ end
+ end
end
- it "routes to short branch for short strings" do
- result = gate.run("hi")
+ it "uses condition and branch DSL" do
+ short_cls = short_step_class
+ long_cls = long_step_class
- expect(result).to eq(short_branch)
+ gate_class = Class.new(described_class) do
+ condition { |input| input.length < 5 ? :short : :long }
+ branch :short, short_cls
+ branch :long, long_cls
+ end
+
+ gate = gate_class.new("DSLGate")
+ expect(gate.run("hi")).to eq("quick: hi")
+ expect(gate.run("longstring")).to eq("deep: longstring")
end
end
context "with complex condition logic" do
- let(:condition) do
- lambda do |input|
- case input
- when 0..10 then "low"
- when 11..50 then "medium"
- else "high"
- end
- end
+ let(:low_step) do
+ Class.new(MARS::Runnable) { def run(input) = "low:#{input}" }.new
end
- let(:low_branch) { instance_spy(MARS::Runnable) }
- let(:medium_branch) { instance_spy(MARS::Runnable) }
- let(:high_branch) { instance_spy(MARS::Runnable) }
- let(:branches) { { "low" => low_branch, "medium" => medium_branch, "high" => high_branch } }
+ let(:medium_step) do
+ Class.new(MARS::Runnable) { def run(input) = "med:#{input}" }.new
+ end
- it "routes to low branch" do
- result = gate.run(5)
+ let(:high_step) do
+ Class.new(MARS::Runnable) { def run(input) = "high:#{input}" }.new
+ end
- expect(result).to eq(low_branch)
+ let(:gate) do
+ described_class.new(
+ "SeverityGate",
+ condition: lambda { |input|
+ case input
+ when 0..10 then :low
+ when 11..50 then :medium
+ else :high
+ end
+ },
+ branches: { low: low_step, medium: medium_step, high: high_step }
+ )
end
- it "routes to medium branch" do
- result = gate.run(25)
+ it "routes to low branch" do
+ expect(gate.run(5)).to eq("low:5")
+ end
- expect(result).to eq(medium_branch)
+ it "routes to medium branch" do
+ expect(gate.run(25)).to eq("med:25")
end
it "routes to high branch" do
- result = gate.run(100)
-
- expect(result).to eq(high_branch)
+ expect(gate.run(100)).to eq("high:100")
end
end
end
diff --git a/spec/mars/hooks_spec.rb b/spec/mars/hooks_spec.rb
new file mode 100644
index 0000000..d304a7c
--- /dev/null
+++ b/spec/mars/hooks_spec.rb
@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::Hooks do
+ let(:hookable_class) do
+ Class.new do
+ include MARS::Hooks
+
+ attr_reader :name
+
+ def initialize(name)
+ @name = name
+ end
+ end
+ end
+
+ describe ".before_run" do
+ it "registers a before hook" do
+ hookable_class.before_run { |_ctx, _step| "before" }
+
+ expect(hookable_class.before_run_hooks.size).to eq(1)
+ end
+ end
+
+ describe ".after_run" do
+ it "registers an after hook" do
+ hookable_class.after_run { |_ctx, _result, _step| "after" }
+
+ expect(hookable_class.after_run_hooks.size).to eq(1)
+ end
+ end
+
+ describe "#run_before_hooks" do
+ it "calls all before hooks with context and step" do
+ calls = []
+ hookable_class.before_run { |ctx, step| calls << [ctx, step.name] }
+
+ step = hookable_class.new("test_step")
+ context = MARS::ExecutionContext.new(input: "input")
+ step.run_before_hooks(context)
+
+ expect(calls).to eq([[context, "test_step"]])
+ end
+
+ it "calls hooks in registration order" do
+ order = []
+ hookable_class.before_run { |_ctx, _step| order << :first }
+ hookable_class.before_run { |_ctx, _step| order << :second }
+
+ step = hookable_class.new("test_step")
+ step.run_before_hooks(MARS::ExecutionContext.new(input: "input"))
+
+ expect(order).to eq(%i[first second])
+ end
+ end
+
+ describe "#run_after_hooks" do
+ it "calls all after hooks with context, result, and step" do
+ calls = []
+ hookable_class.after_run { |ctx, result, step| calls << [ctx, result, step.name] }
+
+ step = hookable_class.new("test_step")
+ context = MARS::ExecutionContext.new(input: "input")
+ step.run_after_hooks(context, "the result")
+
+ expect(calls).to eq([[context, "the result", "test_step"]])
+ end
+ end
+
+ describe "hook isolation between classes" do
+ it "does not share hooks between different classes" do
+ class_a = Class.new { include MARS::Hooks }
+ class_b = Class.new { include MARS::Hooks }
+
+ class_a.before_run { "a" }
+
+ expect(class_a.before_run_hooks.size).to eq(1)
+ expect(class_b.before_run_hooks.size).to eq(0)
+ end
+ end
+end
diff --git a/spec/mars/rendering/graph_spec.rb b/spec/mars/rendering/graph_spec.rb
new file mode 100644
index 0000000..17b1bfe
--- /dev/null
+++ b/spec/mars/rendering/graph_spec.rb
@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::Rendering::Graph do
+ describe "default Runnable rendering" do
+ it "renders any Runnable subclass as a box node" do
+ step_class = Class.new(MARS::Runnable) do
+ def run(input)
+ input
+ end
+ end
+
+ step = step_class.new(name: "custom_step")
+ graph, nodes, _subgraphs = step.build_graph
+
+ expect(nodes).to have_key("custom_step")
+ expect(nodes["custom_step"].type).to eq(MARS::Rendering::Graph::Node::STEP)
+ expect(graph).to have_key("in")
+ end
+ end
+
+ describe "AgentStep rendering" do
+ it "renders as a step node" do
+ mock_agent = Class.new
+ step_class = Class.new(MARS::AgentStep) do
+ agent mock_agent
+ end
+
+ step = step_class.new(name: "my_agent")
+ _graph, nodes, _subgraphs = step.build_graph
+
+ expect(nodes["my_agent"].type).to eq(MARS::Rendering::Graph::Node::STEP)
+ end
+ end
+end
diff --git a/spec/mars/rendering/html_spec.rb b/spec/mars/rendering/html_spec.rb
new file mode 100644
index 0000000..d09b9a9
--- /dev/null
+++ b/spec/mars/rendering/html_spec.rb
@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "tmpdir"
+
+RSpec.describe MARS::Rendering::Html do
+ let(:step_class) do
+ Class.new(MARS::Runnable) do
+ def run(input)
+ input
+ end
+ end
+ end
+
+ let(:workflow) do
+ step_a = step_class.new(name: "step_a")
+ step_b = step_class.new(name: "step_b")
+ MARS::Workflows::Sequential.new("TestPipeline", steps: [step_a, step_b])
+ end
+
+ describe "#render" do
+ it "returns a self-contained HTML string" do
+ html = described_class.new(workflow).render
+
+ expect(html).to include("")
+ expect(html).to include("beautiful-mermaid")
+ expect(html).to include("renderMermaidSVG")
+ expect(html).to include("TestPipeline")
+ expect(html).to include("flowchart LR")
+ end
+
+ it "accepts a custom direction" do
+ html = described_class.new(workflow).render(direction: "TD")
+
+ expect(html).to include("flowchart TD")
+ end
+
+ it "accepts a custom title" do
+ html = described_class.new(workflow).render(title: "My Workflow")
+
+ expect(html).to include("My Workflow")
+ end
+
+ it "accepts theme options" do
+ html = described_class.new(workflow).render(theme: { bg: "#1a1b26", fg: "#a9b1d6" })
+
+ expect(html).to include("bg: '#1a1b26'")
+ expect(html).to include("fg: '#a9b1d6'")
+ end
+
+ it "escapes HTML in the title" do
+ html = described_class.new(workflow).render(title: "")
+
+ expect(html).not_to include("