Update complex examples

examples/complex_llm_workflow/generator.rb

@@ -37,41 +37,42 @@ def execute(latitude:, longitude:)
   end
 end
 
-# Define LLMs
-class Agent1 < MARS::AgentStep
-  def system_prompt
-    "You are a helpful assistant that can answer questions.
+class ExampleAgent < RubyLLM::Agent
+  instructions "You are a helpful assistant that can answer questions.
      When asked about a country, only answer with its name."
-  end
 end
 
-class Agent2 < MARS::AgentStep
-  def system_prompt
-    "You are a helpful assistant that can answer questions and help with tasks.
+class Agent1 < MARS::AgentStep
+  agent ExampleAgent
+end
+
+class FoodAgent < RubyLLM::Agent
+  instructions "You are a helpful assistant that can answer questions and help with tasks.
      Return information about the typical food of the country."
-  end
 end
 
-class Agent3 < MARS::AgentStep
-  def system_prompt
-    "You are a helpful assistant that can answer questions and help with tasks.
+class Agent2 < MARS::AgentStep
+  agent FoodAgent
+end
+
+class SportsAgent < RubyLLM::Agent
+  instructions "You are a helpful assistant that can answer questions and help with tasks.
      Return information about the popular sports of the country."
-  end
+  schema SportsSchema
+end
 
-  def schema
-    SportsSchema.new
-  end
+class Agent3 < MARS::AgentStep
+  agent SportsAgent
 end
 
-class Agent4 < MARS::AgentStep
-  def system_prompt
-    "You are a helpful assistant that can answer questions and help with tasks.
+class WeatherAgent < RubyLLM::Agent
+  instructions "You are a helpful assistant that can answer questions and help with tasks.
      Return the current weather of the country's capital."
-  end
+  tools Weather
+end
 
-  def tools
-    [Weather.new]
-  end
+class Agent4 < MARS::AgentStep
+  agent WeatherAgent
 end
 
 # Create the LLMs

examples/complex_workflow/diagram.md

@@ -2,46 +2,46 @@
 flowchart LR
 in((In))
 out((Out))
-agent1[agent1]
+step1[step1]
 gate{Gate}
-agent4[agent4]
+step4[step4]
 parallel_workflow_aggregator[Parallel workflow Aggregator]
-agent2[agent2]
-agent3[agent3]
+step2[step2]
+step3[step3]
 parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator]
-agent5[agent5]
+step5[step5]
 subgraph parallel_workflow_2["Parallel workflow 2"]
   sequential_workflow
-  agent5
+  step5
 end
 subgraph parallel_workflow["Parallel workflow"]
-  agent2
-  agent3
+  step2
+  step3
 end
 subgraph sequential_workflow["Sequential workflow"]
-  agent4
+  step4
   parallel_workflow
   parallel_workflow_aggregator
 end
 subgraph main_pipeline["Main Pipeline"]
-  agent1
+  step1
   gate
   parallel_workflow_aggregator
   parallel_workflow_2
   parallel_workflow_2_aggregator
 end
-in --> agent1
-agent1 --> gate
-gate -->|warning| agent4
-gate -->|error| agent2
-gate -->|error| agent3
-gate --> agent4
-gate --> agent5
-agent4 --> agent2
-agent4 --> agent3
-agent2 --> parallel_workflow_aggregator
+in --> step1
+step1 --> gate
+gate -->|warning| step4
+gate -->|error| step2
+gate -->|error| step3
+gate --> step4
+gate --> step5
+step4 --> step2
+step4 --> step3
+step2 --> parallel_workflow_aggregator
 parallel_workflow_aggregator --> parallel_workflow_2_aggregator
-agent3 --> parallel_workflow_aggregator
+step3 --> parallel_workflow_aggregator
 parallel_workflow_2_aggregator --> out
-agent5 --> parallel_workflow_2_aggregator
+step5 --> parallel_workflow_2_aggregator
 ```

examples/complex_workflow/generator.rb

@@ -3,50 +3,91 @@
 
 require_relative "../../lib/mars"
 
-# Define LLMs
-class Agent1 < MARS::AgentStep
+class Step1Formatter < MARS::Formatter
+  def format_output(context)
+    context + ['formatted']
+  end
 end
 
-class Agent2 < MARS::AgentStep
+class Step1 < MARS::Runnable
+  after_run do |context, result|
+    puts "after run from Step1 #{result}"
+  end
+  formatter Step1Formatter
+
+  def run(context)
+    context.current_input + ['step1']
+  end
+end
+
+class Step2 < MARS::Runnable
+  after_run do |context, result|
+    puts "after run from Step2 #{result}"
+  end
+
+  def run(context)
+    context.current_input + ['step2']
+  end
 end
 
-class Agent3 < MARS::AgentStep
+class Step3 < MARS::Runnable
+  after_run do |context, result|
+   puts "after run from Step3 #{result}"
+  end
+
+  def run(context)
+    context.current_input + ['step3']
+  end
 end
 
-class Agent4 < MARS::AgentStep
+class Step4 < MARS::Runnable
+  after_run do |context, result|
+    puts "after run from Step4 #{result}"
+  end
+
+  def run(context)
+    context.current_input + ['step4']
+  end
 end
 
-class Agent5 < MARS::AgentStep
+class Step5 < MARS::Runnable
+  after_run do |context, result|
+    puts "after run from Step5 #{result}"
+  end
+
+  def run(context)
+    context.current_input + ['step5']
+  end
 end
 
-# Create the LLMs
-llm1 = Agent1.new
-llm2 = Agent2.new
-llm3 = Agent3.new
-llm4 = Agent4.new
-llm5 = Agent5.new
+# Create the Steps
+step1 = Step1.new
+step2 = Step2.new
+step3 = Step3.new
+step4 = Step4.new
+step5 = Step5.new
 
-# Create a parallel workflow (LLM 2 x LLM 3)
+# Create a parallel workflow (STEP 2 x STEP 3)
 parallel_workflow = MARS::Workflows::Parallel.new(
   "Parallel workflow",
-  steps: [llm2, llm3]
+  steps: [step2, step3]
 )
 
 # Create a sequential workflow (Parallel workflow -> LLM 4)
 sequential_workflow = MARS::Workflows::Sequential.new(
   "Sequential workflow",
-  steps: [llm4, parallel_workflow]
+  steps: [step4, parallel_workflow]
 )
 
 # Create a parallel workflow (Sequential workflow x LLM 5)
 parallel_workflow2 = MARS::Workflows::Parallel.new(
   "Parallel workflow 2",
-  steps: [sequential_workflow, llm5]
+  steps: [sequential_workflow, step5]
 )
 
 # Create the gate that decides between exit or continue
 gate = MARS::Gate.new(
-  check: ->(input) { input[:result] },
+  check: ->(input) { :default if input == ["start", "step1", "formatted"] },
   fallbacks: {
     warning: sequential_workflow,
     error: parallel_workflow
@@ -56,9 +97,11 @@ class Agent5 < MARS::AgentStep
 # Create the main workflow: LLM 1 -> Gate
 main_workflow = MARS::Workflows::Sequential.new(
   "Main Pipeline",
-  steps: [llm1, gate, parallel_workflow2]
+  steps: [step1, gate, parallel_workflow2]
 )
 
+main_workflow.run(["start"])
+
 # Generate and save the diagram
 diagram = MARS::Rendering::Mermaid.new(main_workflow).render
 File.write("examples/complex_workflow/diagram.md", diagram)

lib/mars/gate.rb

@@ -30,7 +30,7 @@ def run(context)
       input = context.current_input
       result = check.call(input)
 
-      return input unless result
+      return input if result.nil? || result == :default
 
       branch = fallbacks[result]
       raise ArgumentError, "No fallback registered for #{result.inspect}" unless branch