From 23bde3f76a6ec7c36ab1f80311cb9c446e8fb836 Mon Sep 17 00:00:00 2001
From: Santiago Diaz <santiago.d.diaz@gmail.com>
Date: Fri, 31 Oct 2025 16:34:16 -0300
Subject: [PATCH] Add complex example with LLM usage

---
 examples/complex_llm_workflow/generator.rb | 42 ++++++++++++++++++++++
 lib/mars.rb                                |  1 +
 lib/mars/agent.rb                          | 12 ++++---
 lib/mars/gate.rb                           |  2 +-
 4 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100755 examples/complex_llm_workflow/generator.rb

diff --git a/examples/complex_llm_workflow/generator.rb b/examples/complex_llm_workflow/generator.rb
new file mode 100755
index 0000000..d96acd4
--- /dev/null
+++ b/examples/complex_llm_workflow/generator.rb
@@ -0,0 +1,42 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require_relative "../../lib/mars"
+
+RubyLLM.configure do |config|
+  config.openai_api_key = ENV.fetch("OPENAI_API_KEY", nil)
+end
+
+# Create the LLMs
+llm1 = Mars::Agent.new(
+  name: "LLM 1", options: { model: "gpt-4o" },
+  instructions: "You are a helpful assistant that can answer questions and help with tasks. Only answer with the result"
+)
+
+llm2 = Mars::Agent.new(name: "LLM 2", options: { model: "gpt-4o" },
+                       instructions: "You are a helpful assistant that can answer questions and help with tasks.
+                       Return information about the typical food of the country.")
+
+llm3 = Mars::Agent.new(name: "LLM 3", options: { model: "gpt-4o" },
+                       instructions: "You are a helpful assistant that can answer questions and help with tasks.
+                       Return information about the popular sports of the country.")
+
+parallel_workflow = Mars::Workflows::Parallel.new(
+  "Parallel workflow",
+  steps: [llm2, llm3]
+)
+
+gate = Mars::Gate.new(
+  name: "Gate",
+  condition: ->(input) { input.split.length < 10 ? :success : :error },
+  branches: {
+    success: parallel_workflow
+  }
+)
+
+sequential_workflow = Mars::Workflows::Sequential.new(
+  "Sequential workflow",
+  steps: [llm1, gate]
+)
+
+puts sequential_workflow.run("Which is the last country to declare independence?")
diff --git a/lib/mars.rb b/lib/mars.rb
index 06cce5d..68e23d7 100644
--- a/lib/mars.rb
+++ b/lib/mars.rb
@@ -2,6 +2,7 @@
 
 require "zeitwerk"
 require "async"
+require "ruby_llm"
 
 loader = Zeitwerk::Loader.for_gem
 loader.setup
diff --git a/lib/mars/agent.rb b/lib/mars/agent.rb
index 1d39a67..66eaa49 100644
--- a/lib/mars/agent.rb
+++ b/lib/mars/agent.rb
@@ -4,23 +4,27 @@ module Mars
   class Agent < Runnable
     attr_reader :name
 
-    def initialize(name:, options: {}, tools: [], schema: nil)
+    def initialize(name:, options: {}, tools: [], schema: nil, instructions: nil)
       @name = name
       @tools = Array(tools)
       @schema = schema
       @options = options
+      @instructions = instructions
     end
 
     def run(input)
-      chat.ask(input)
+      chat.ask(input).content
     end
 
     private
 
-    attr_reader :tools, :schema, :options
+    attr_reader :tools, :schema, :options, :instructions
 
     def chat
-      @chat ||= RubyLLM::Chat.new(**options).with_tools(tools).with_schema(schema)
+      @chat ||= RubyLLM::Chat.new(**options)
+                             .with_instructions(instructions)
+                             .with_tools(*tools)
+                             .with_schema(schema)
     end
   end
 end
diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb
index f838912..5ac6aa1 100644
--- a/lib/mars/gate.rb
+++ b/lib/mars/gate.rb
@@ -13,7 +13,7 @@ def initialize(name:, condition:, branches:)
     def run(input)
       result = condition.call(input)
 
-      branches[result].run(input)
+      branches[result]&.run(input) || input
     end
 
     private