brute 0.4.1 → 1.0.1

data/lib/brute/middleware/otel/tool_calls.rb

@@ -1,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -20,7 +18,7 @@ module Brute
 
           span = env[:span]
           if span
-            functions = env[:context].functions
+            functions = env[:pending_functions]
             if functions && !functions.empty?
               span.set_attribute("brute.tool_calls.count", functions.size)
 
@@ -43,110 +41,28 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::OTel::ToolCalls do
-    let(:response) { MockResponse.new(content: "here's my plan") }
-    let(:inner_app) { ->(_env) { response } }
-    let(:middleware) { described_class.new(inner_app) }
-
-    it "passes the response through unchanged" do
-      env = build_env
-      result = middleware.call(env)
-      expect(result).to eq(response)
-    end
-
-    context "when env[:span] is nil" do
-      it "passes through without error even with pending functions" do
-        ctx = build_env[:context]
-        fn = double("function", name: "fs_read", id: "tc_001", arguments: { "path" => "/tmp" })
-        allow(ctx).to receive(:functions).and_return([fn])
-
-        env = build_env(context: ctx)
-        result = middleware.call(env)
-        expect(result).to eq(response)
-      end
-    end
-
-    context "when env[:span] is present" do
-      let(:span) { mock_span }
-
-      it "does nothing when there are no pending functions" do
-        ctx = build_env[:context]
-        allow(ctx).to receive(:functions).and_return([])
-
-        env = build_env(context: ctx, span: span)
-        middleware.call(env)
-
-        expect(span).not_to have_received(:add_event)
-        expect(span).not_to have_received(:set_attribute)
-      end
-
-      it "does nothing when functions is nil" do
-        ctx = build_env[:context]
-        allow(ctx).to receive(:functions).and_return(nil)
-
-        env = build_env(context: ctx, span: span)
-        middleware.call(env)
+test do
+  require_relative "../../../../spec/support/mock_provider"
+  require_relative "../../../../spec/support/mock_response"
 
-        expect(span).not_to have_received(:add_event)
-      end
-
-      it "records a tool_call event per pending function" do
-        ctx = build_env[:context]
-        fn1 = double("function", name: "fs_read", id: "tc_001", arguments: { "path" => "/src/main.rb" })
-        fn2 = double("function", name: "shell", id: "tc_002", arguments: { "command" => "rspec" })
-        allow(ctx).to receive(:functions).and_return([fn1, fn2])
-
-        env = build_env(context: ctx, span: span)
-        middleware.call(env)
-
-        expect(span).to have_received(:set_attribute).with("brute.tool_calls.count", 2)
-        expect(span).to have_received(:add_event).with(
-          "tool_call",
-          attributes: hash_including(
-            "tool.name" => "fs_read",
-            "tool.id" => "tc_001"
-          )
-        )
-        expect(span).to have_received(:add_event).with(
-          "tool_call",
-          attributes: hash_including(
-            "tool.name" => "shell",
-            "tool.id" => "tc_002"
-          )
-        )
-      end
-
-      it "serializes arguments as JSON" do
-        ctx = build_env[:context]
-        args = { "path" => "/tmp/test.rb", "content" => "puts 'hi'" }
-        fn = double("function", name: "fs_write", id: "tc_003", arguments: args)
-        allow(ctx).to receive(:functions).and_return([fn])
-
-        env = build_env(context: ctx, span: span)
-        middleware.call(env)
-
-        expect(span).to have_received(:add_event).with(
-          "tool_call",
-          attributes: hash_including("tool.arguments" => args.to_json)
-        )
-      end
-
-      it "handles nil arguments" do
-        ctx = build_env[:context]
-        fn = double("function", name: "todo_read", id: "tc_004", arguments: nil)
-        allow(ctx).to receive(:functions).and_return([fn])
+  def build_env(**overrides)
+    { provider: MockProvider.new, model: nil, input: "test prompt", tools: [],
+      messages: [], stream: nil, params: {}, metadata: {}, callbacks: {},
+      tool_results: nil, streaming: false, should_exit: nil, pending_functions: [] }.merge(overrides)
+  end
 
-        env = build_env(context: ctx, span: span)
-        middleware.call(env)
+  it "passes the response through unchanged" do
+    response = MockResponse.new(content: "here's my plan")
+    middleware = Brute::Middleware::OTel::ToolCalls.new(->(_env) { response })
+    result = middleware.call(build_env)
+    result.should == response
+  end
 
-        expect(span).to have_received(:add_event).with(
-          "tool_call",
-          attributes: { "tool.name" => "todo_read", "tool.id" => "tc_004" }
-        )
-      end
-    end
+  it "passes through without error when span is nil with pending functions" do
+    response = MockResponse.new(content: "here's my plan")
+    fn = Struct.new(:name, :id, :arguments, keyword_init: true).new(name: "fs_read", id: "tc_001", arguments: { "path" => "/tmp" })
+    middleware = Brute::Middleware::OTel::ToolCalls.new(->(_env) { response })
+    result = middleware.call(build_env(pending_functions: [fn]))
+    result.should == response
   end
 end

data/lib/brute/middleware/otel/tool_results.rb

@@ -1,16 +1,14 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
     module OTel
       # Records tool results being sent back to the LLM as span events.
       #
-      # Runs PRE-call: when env[:tool_results] is present, the orchestrator
+      # Runs PRE-call: when env[:tool_results] is present, the agent loop
       # is sending tool execution results back to the LLM. Each result gets
       # a span event with the tool name and success/error status.
       #
@@ -41,89 +39,27 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::OTel::ToolResults do
-    let(:response) { MockResponse.new(content: "processed") }
-    let(:inner_app) { ->(_env) { response } }
-    let(:middleware) { described_class.new(inner_app) }
-
-    it "passes the response through unchanged" do
-      env = build_env
-      result = middleware.call(env)
-      expect(result).to eq(response)
-    end
-
-    context "when env[:span] is nil" do
-      it "passes through without error" do
-        results = [["fs_read", { content: "data" }]]
-        env = build_env(tool_results: results)
+test do
+  require_relative "../../../../spec/support/mock_provider"
+  require_relative "../../../../spec/support/mock_response"
 
-        result = middleware.call(env)
-        expect(result).to eq(response)
-      end
-    end
-
-    context "when env[:span] is present" do
-      let(:span) { mock_span }
-
-      it "does nothing when tool_results is nil" do
-        env = build_env(span: span, tool_results: nil)
-        middleware.call(env)
-
-        expect(span).not_to have_received(:add_event)
-        expect(span).not_to have_received(:set_attribute)
-      end
-
-      it "records a tool_result event per result" do
-        results = [
-          ["fs_read", { content: "file data" }],
-          ["shell", { output: "ok" }],
-        ]
-        env = build_env(span: span, tool_results: results)
-        middleware.call(env)
-
-        expect(span).to have_received(:set_attribute).with("brute.tool_results.count", 2)
-        expect(span).to have_received(:add_event).with(
-          "tool_result",
-          attributes: hash_including("tool.name" => "fs_read", "tool.status" => "ok")
-        )
-        expect(span).to have_received(:add_event).with(
-          "tool_result",
-          attributes: hash_including("tool.name" => "shell", "tool.status" => "ok")
-        )
-      end
-
-      it "records error status and message for failed tool results" do
-        results = [
-          ["fs_read", { error: "not found" }],
-        ]
-        env = build_env(span: span, tool_results: results)
-        middleware.call(env)
-
-        expect(span).to have_received(:add_event).with(
-          "tool_result",
-          attributes: hash_including(
-            "tool.name" => "fs_read",
-            "tool.status" => "error",
-            "tool.error" => "not found"
-          )
-        )
-      end
+  def build_env(**overrides)
+    { provider: MockProvider.new, model: nil, input: "test prompt", tools: [],
+      messages: [], stream: nil, params: {}, metadata: {}, callbacks: {},
+      tool_results: nil, streaming: false, should_exit: nil, pending_functions: [] }.merge(overrides)
+  end
 
-      it "handles a mix of successful and failed results" do
-        results = [
-          ["fs_read", { content: "ok" }],
-          ["shell", { error: "exit code 1" }],
-          ["fs_write", { success: true }],
-        ]
-        env = build_env(span: span, tool_results: results)
-        middleware.call(env)
+  it "passes the response through unchanged" do
+    response = MockResponse.new(content: "processed")
+    middleware = Brute::Middleware::OTel::ToolResults.new(->(_env) { response })
+    result = middleware.call(build_env)
+    result.should == response
+  end
 
-        expect(span).to have_received(:set_attribute).with("brute.tool_results.count", 3)
-        expect(span).to have_received(:add_event).exactly(3).times
-      end
-    end
+  it "passes through without error when span is nil" do
+    response = MockResponse.new(content: "processed")
+    middleware = Brute::Middleware::OTel::ToolResults.new(->(_env) { response })
+    result = middleware.call(build_env(tool_results: [["fs_read", { content: "data" }]]))
+    result.should == response
   end
 end

data/lib/brute/middleware/reasoning_normalizer.rb

@@ -1,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -102,116 +100,93 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::ReasoningNormalizer do
-    let(:response) { MockResponse.new(content: "reasoned response") }
-    let(:inner_app) { ->(_env) { response } }
-
-    # Build a provider whose class name contains the given string.
-    def make_provider(type_name)
-      klass = Class.new do
-        define_method(:name) { :mock }
-        define_method(:default_model) { "mock-model" }
-        define_method(:user_role) { :user }
-        define_method(:system_role) { :system }
-        define_method(:assistant_role) { :assistant }
-        define_method(:tool_role) { :tool }
-        define_method(:tracer) { nil }
-        define_method(:tracer=) { |*| }
-        define_method(:complete) { |*_args, **_kw| MockResponse.new(content: "ok") }
-      end
-      # Override class name to trigger provider_type detection
-      klass.define_method(:class) do
-        c = super()
-        name_str = "LLM::#{type_name}"
-        c.define_singleton_method(:name) { name_str }
-        c
-      end
-      klass.new
-    end
-
-    context "with Anthropic provider and budget_tokens" do
-      it "injects thinking param into env[:params]" do
-        provider = make_provider("Anthropic")
-        middleware = described_class.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: true)
-        env = build_env(provider: provider, params: {})
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-        middleware.call(env)
+  def build_env(**overrides)
+    { provider: MockProvider.new, model: nil, input: "test prompt", tools: [],
+      messages: [], stream: nil, params: {}, metadata: {}, callbacks: {},
+      tool_results: nil, streaming: false, should_exit: nil, pending_functions: [] }.merge(overrides)
+  end
 
-        expect(env[:params][:thinking]).to eq({ type: "enabled", budget_tokens: 8000 })
-      end
+  def make_provider(type_name)
+    klass = Class.new do
+      define_method(:name) { :mock }
+      define_method(:default_model) { "mock-model" }
+      define_method(:user_role) { :user }
+      define_method(:system_role) { :system }
+      define_method(:assistant_role) { :assistant }
+      define_method(:tool_role) { :tool }
+      define_method(:tracer) { nil }
+      define_method(:tracer=) { |*| }
+      define_method(:complete) { |*_args, **_kw| MockResponse.new(content: "ok") }
     end
-
-    context "with Anthropic provider without budget_tokens" do
-      it "does not inject thinking param" do
-        provider = make_provider("Anthropic")
-        middleware = described_class.new(inner_app, model_id: "claude-4", enabled: true)
-        env = build_env(provider: provider, params: {})
-
-        middleware.call(env)
-
-        expect(env[:params][:thinking]).to be_nil
-      end
+    klass.define_method(:class) do
+      c = super()
+      name_str = "LLM::#{type_name}"
+      c.define_singleton_method(:name) { name_str }
+      c
     end
+    klass.new
+  end
 
-    context "with OpenAI provider" do
-      it "injects reasoning_effort based on effort level" do
-        provider = make_provider("OpenAI")
-        middleware = described_class.new(inner_app, model_id: "o3", effort: :high, enabled: true)
-        env = build_env(provider: provider, params: {})
-
-        middleware.call(env)
-
-        expect(env[:params][:reasoning_effort]).to eq("high")
-      end
-
-      it "maps effort levels correctly" do
-        provider = make_provider("OpenAI")
-
-        { low: "low", medium: "medium", high: "high", minimal: "low", max: "high" }.each do |effort, expected|
-          middleware = described_class.new(inner_app, model_id: "o3", effort: effort, enabled: true)
-          env = build_env(provider: provider, params: {})
-          middleware.call(env)
-          expect(env[:params][:reasoning_effort]).to eq(expected), "Expected effort #{effort} to map to #{expected}"
-        end
-      end
-    end
-
-    context "with unknown provider" do
-      it "does not inject any reasoning params" do
-        provider = make_provider("Mistral")
-        middleware = described_class.new(inner_app, model_id: "mistral-large", enabled: true)
-        env = build_env(provider: provider, params: {})
-
-        middleware.call(env)
+  inner_app = ->(_env) { MockResponse.new(content: "reasoned response") }
 
-        expect(env[:params]).to eq({})
-      end
-    end
+  it "injects thinking param for Anthropic with budget_tokens" do
+    provider = make_provider("Anthropic")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:thinking].should == { type: "enabled", budget_tokens: 8000 }
+  end
 
-    context "when disabled" do
-      it "does not inject reasoning params" do
-        provider = make_provider("Anthropic")
-        middleware = described_class.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: false)
-        env = build_env(provider: provider, params: {})
+  it "does not inject thinking param for Anthropic without budget_tokens" do
+    provider = make_provider("Anthropic")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "claude-4", enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:thinking].should.be.nil
+  end
 
-        middleware.call(env)
+  it "injects reasoning_effort for OpenAI" do
+    provider = make_provider("OpenAI")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "o3", effort: :high, enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:reasoning_effort].should == "high"
+  end
 
-        expect(env[:params]).to eq({})
-      end
-    end
+  it "maps low effort correctly for OpenAI" do
+    provider = make_provider("OpenAI")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "o3", effort: :low, enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:reasoning_effort].should == "low"
+  end
 
-    it "allows model_id to be updated mid-session" do
-      middleware = described_class.new(inner_app, model_id: "old-model", enabled: true)
-      middleware.model_id = "new-model"
+  it "does not inject params for unknown provider" do
+    provider = make_provider("Mistral")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "mistral-large", enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params].should == {}
+  end
 
-      provider = make_provider("OpenAI")
-      env = build_env(provider: provider, params: {})
-      middleware.call(env)
+  it "does not inject params when disabled" do
+    provider = make_provider("Anthropic")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: false)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params].should == {}
+  end
 
-      expect(env[:params][:reasoning_effort]).not_to be_nil
-    end
+  it "allows model_id to be updated mid-session" do
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "old", enabled: true)
+    middleware.model_id = "new"
+    provider = make_provider("OpenAI")
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:reasoning_effort].should.not.be.nil
   end
 end