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/rendering/graph.rb b/lib/mars/rendering/graph.rb index 1a73d66..b2ab350 100644 --- a/lib/mars/rendering/graph.rb +++ b/lib/mars/rendering/graph.rb @@ -4,7 +4,7 @@ module MARS module Rendering module Graph def self.include_extensions - MARS::Agent.include(Agent) + 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/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