diff --git a/examples/complex_llm_workflow/generator.rb b/examples/complex_llm_workflow/generator.rb index 6bc1436..46d5e0b 100755 --- a/examples/complex_llm_workflow/generator.rb +++ b/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 diff --git a/examples/complex_workflow/diagram.md b/examples/complex_workflow/diagram.md index 51cf240..7dcbfa7 100644 --- a/examples/complex_workflow/diagram.md +++ b/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 ``` diff --git a/examples/complex_workflow/generator.rb b/examples/complex_workflow/generator.rb index 1247e6d..4277a29 100755 --- a/examples/complex_workflow/generator.rb +++ b/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) diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb index 195716b..5418d73 100644 --- a/lib/mars/gate.rb +++ b/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