Fix examples (#33)

* Fix simple_workflow example

* Fix parallel_workflow example

* Fix complex_workflow example

* Fix complex_llm_workflow example

Author: santib

examples/complex_llm_workflow/diagram.md

@@ -2,12 +2,12 @@
 flowchart LR
 in((In))
 out((Out))
-agent1[Agent1]
+agent1[agent1]
 gate{Gate}
 parallel_workflow_aggregator[Parallel workflow Aggregator]
-agent2[Agent2]
-agent3[Agent3]
-agent4[Agent4]
+agent2[agent2]
+agent3[agent3]
+agent4[agent4]
 subgraph parallel_workflow["Parallel workflow"]
   agent2
   agent3

examples/complex_llm_workflow/generator.rb

@@ -38,21 +38,21 @@ def execute(latitude:, longitude:)
 end
 
 # Define LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
   def system_prompt
     "You are a helpful assistant that can answer questions.
      When asked about a country, only answer with its name."
   end
 end
 
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
   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
 end
 
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
   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."
@@ -63,7 +63,7 @@ def schema
   end
 end
 
-class Agent4 < MARS::Agent
+class Agent4 < MARS::AgentStep
   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."
@@ -75,10 +75,10 @@ def tools
 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" })
+llm1 = Agent1.new
+llm2 = Agent2.new
+llm3 = Agent3.new
+llm4 = Agent4.new
 
 parallel_workflow = MARS::Workflows::Parallel.new(
   "Parallel workflow",

examples/complex_workflow/diagram.md

@@ -2,14 +2,14 @@
 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]
+agent5[agent5]
 subgraph parallel_workflow_2["Parallel workflow 2"]
   sequential_workflow
   agent5

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

examples/parallel_workflow/diagram.md

@@ -3,9 +3,9 @@ flowchart LR
 in((In))
 out((Out))
 aggregator[Aggregator]
-agent1[Agent1]
-agent2[Agent2]
-agent3[Agent3]
+agent1[agent1]
+agent2[agent2]
+agent3[agent3]
 subgraph parallel_workflow["Parallel workflow"]
   agent1
   agent2

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

examples/simple_workflow/diagram.md

@@ -2,14 +2,25 @@
 flowchart LR
 in((In))
 out((Out))
-agent1[Agent1]
+agent1[agent1]
 gate{Gate}
-agent2[Agent2]
-agent3[Agent3]
+agent4[agent4]
+agent2[agent2]
+agent3[agent3]
+subgraph failure_workflow["Failure workflow"]
+  agent4
+end
+subgraph main_pipeline["Main Pipeline"]
+  agent1
+  gate
+  agent4
+  agent2
+  agent3
+end
 in --> agent1
 agent1 --> gate
-gate -->|success| agent2
-gate -->|default| out
+gate -->|failure| agent4
+gate --> agent2
 agent2 --> agent3
 agent3 --> out
 ```

examples/simple_workflow/generator.rb

@@ -4,38 +4,42 @@
 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
+
+class Agent4 < MARS::AgentStep
 end
 
 # Create the LLMs
 llm1 = Agent1.new
 llm2 = Agent2.new
 llm3 = Agent3.new
+llm4 = Agent4.new
 
-# Create the success workflow (LLM 2 -> LLM 3)
-success_workflow = MARS::Workflows::Sequential.new(
-  "Success workflow",
-  steps: [llm2, llm3]
+# 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: {
-    success: success_workflow
+    failure: failure_workflow
   }
 )
 
 # Create the main workflow: LLM 1 -> Gate
 main_workflow = MARS::Workflows::Sequential.new(
   "Main Pipeline",
-  steps: [llm1, gate]
+  steps: [llm1, gate, llm2, llm3]
 )
 
 # Generate and save the diagram

lib/mars.rb

@@ -7,6 +7,7 @@
 
 loader = Zeitwerk::Loader.for_gem
 loader.inflector.inflect("mars" => "MARS")
+loader.ignore("#{__dir__}/mars_rb.rb")
 loader.setup
 
 module MARS